11institutetext: Department of Computer Science, University of Bari Aldo Moro,
Italy 11email<EMAIL_ADDRESS>
22institutetext: NVIDIA AI Technology Center, Bologna, Italy
22email<EMAIL_ADDRESS>
# Label Anything: Multi-Class Few-Shot Semantic Segmentation with Visual
Prompts

Pasquale De Marinis 11 0000-0001-8935-9156 Nicola Fanelli 11
0009-0007-6602-7504 Raffaele Scaringi 11 0000-0001-7512-7661 Emanuele Colonna
11 0009-0009-0932-3424 Giuseppe Fiameni 22 0000-0001-8687-6609 Gennaro Vessio
11 0000-0002-0883-2691 Giovanna Castellano 11 0000-0002-6489-8628

###### Abstract

We present Label Anything, an innovative neural network architecture designed
for few-shot semantic segmentation (FSS) that demonstrates remarkable
generalizability across multiple classes with minimal examples required per
class. Diverging from traditional FSS methods that predominantly rely on masks
for annotating support images, Label Anything introduces varied visual
prompts—points, bounding boxes, and masks—thereby enhancing the framework’s
versatility and adaptability. Unique to our approach, Label Anything is
engineered for end-to-end training across multi-class FSS scenarios,
efficiently learning from diverse support set configurations without
retraining. This approach enables a “universal” application to various FSS
challenges, ranging from $1$-way $1$-shot to complex $N$-way $K$-shot
configurations while remaining agnostic to the specific number of class
examples. This innovative training strategy reduces computational requirements
and substantially improves the model’s adaptability and generalization across
diverse segmentation tasks. Our comprehensive experimental validation,
particularly achieving state-of-the-art results on the COCO-$20^{i}$
benchmark, underscores Label Anything’s robust generalization and flexibility.
The source code is publicly available at:
https://github.com/pasqualedem/LabelAnything.

###### Keywords:

Few-Shot Semantic Segmentation, Visual Prompts, Multi-Class Segmentation,
Class Prototypes, Deep Learning Architectures

Figure 1: Overview of the proposed framework, illustrating the workflow for a
1-shot scenario involving three classes, two support images, and all three
types of prompts. Label Anything leverages a pre-trained ViT to derive image
features. Prompts are fed into the Prompt Encoder, which generates a distinct
prototype for each class. Subsequently, these prototypes inform the Mask
Decoder, facilitating accurate pixel classification across the query image.

##  1 Introduction

The advancement of semantic segmentation has been closely linked to the
availability of extensive labeled datasets, which often pose a significant
challenge in deploying these systems for real-world applications. Typically,
models trained on such datasets excel at recognizing and classifying objects
they were trained on but falter when encountering new concepts or domains,
thus limiting their practical utility. Segmenting objects from previously
unseen classes with few or no labeled examples remains a pivotal challenge in
computer vision.

Few-shot semantic segmentation (FSS) has recently gained attention as the task
of segmenting a query image based on a limited set of examples [26]. While the
majority of research in FSS has concentrated on binary segmentation tasks
(distinguishing background from foreground), there has been limited
exploration of methods capable of segmenting multiple classes simultaneously,
a step toward more flexible and domain-agnostic segmentation [4, 33, 14, 42,
13, 30]. Furthermore, aside from an initial exploration by Rakelly et al.
[24], which considered coordinates as prompts, the field has predominantly
focused on mask annotations as the sole input type. This underscores the
growing need to support diverse input annotations, reducing reliance on
extensive human annotation and facilitating smoother integration with other
models, such as object detectors that generate bounding boxes.

This work presents Label Anything (LA), a novel deep neural network
architecture designed for multi-prompt, multi-class FSS. Our method embraces a
prototype-based learning paradigm, leveraging diverse visual prompts—points,
bounding boxes, and masks—to create a highly flexible and generalizable
framework (Fig. 1 illustrates our approach). This innovation not only
addresses the challenge of segmenting objects from unseen classes with limited
examples but also significantly reduces the annotation burden.

Building upon Segment Anything’s innovative use of diverse prompts for
segmentation [9], our method significantly enriches FSS by enhancing support
set adaptability and extending the framework into multi-class challenges.
Unlike traditional approaches that focus on binary segmentation and adhere to
strict class example limits, LA supports various prompt types and a flexible
number of class examples per support set. This flexibility allows for
capturing complex inter-class relationships from a single image, improving
segmentation precision. Moreover, LA’s architecture is designed to learn from
different support set configurations and prompt types in a single training
cycle, eliminating the need for retraining across different FSS scenarios.
Such versatility ensures that LA can be applied across a wide range of
domains, offering a comprehensive solution that moves beyond the constraints
of conventional FSS models.

Our extensive experimental analysis on the well-known COCO-$20^{i}$ benchmark
[22, 33] reveals that our model achieves state-of-the-art performance, setting
a new standard in the field. The contributions of LA are not merely technical
but also practical as we pave the way for more accessible and efficient
semantic segmentation, especially in scenarios where annotation resources are
scarce.

The rest of this paper is organized as follows. Section 2 reviews related
work. Section 3 details our proposed approach. Section 4 evaluates our model’s
performance. Section 5 concludes the paper with a summary of our findings and
directions for future research.

##  2 Related Work

###  2.1 Few-Shot Semantic Segmentation

Semantic segmentation, which involves assigning a class label to each pixel in
an image, is a fundamental task in computer vision. The introduction of
Convolutional Neural Networks (CNNs), with pioneering architectures like the
Fully-Convolutional Network [17] and U-Net [25], was pivotal in pushing the
field forward. These architectural innovations laid the foundation for
numerous subsequent methods, each tailored to address specific challenges
(e.g., [3, 23]).

Building upon few-shot learning principles [6, 32, 28, 29], few-shot semantic
segmentation has developed as a significant extension of semantic segmentation
to achieve high accuracy with minimal reliance on large, annotated datasets.
FSS involves per-pixel classification in a query image, guided by a “support
set” typically comprising a small number of images and their corresponding
mask annotations. This task hinges on the model’s ability to leverage support
set information for accurate segmentation.

Early attempts in FSS utilized CNNs to synthesize support set information.
This information was integrated into the segmentation process by modifying
weights and biases or merging it with query image features [26, 24]. This
foundational work laid the groundwork for more sophisticated approaches
broadly categorized into prototype-based and affinity-based learning.
Prototype-based methods focus on deriving one or several representative
prototypes for the concept to be segmented, applying these to the query image
through distance metrics or decoder networks [4, 43, 33, 39, 10, 14, 15, 42].
In contrast, affinity learning approaches delve into the pixel-wise
relationships between support and query features [40, 27, 35, 1, 21]. Nearly
all methods in FSS leverage the extensive feature pyramids extracted by CNNs
for query and support images. Only recent works such as FPTrans [41] have
begun to explore the use of the Vision Transformer (ViT) [5] as the encoder.

In recent years, considerable interest has been shown in FSS. However, most
research has focused on a binary paradigm, where images are segmented based on
a single concept specified by the support set. The binary (single-way) and
multi-class FSS settings present distinct challenges. Binary classification
models struggle to adapt to multi-class scenarios without a significant drop
in performance, highlighting the need for specialized solutions.

A limited number of studies have ventured into the complex domain of multi-
class FSS [4, 33, 14, 42, 13, 30]. This involves enabling the support set to
introduce multiple classes for segmentation, requiring the model to
distinguish between $N$ predefined classes in the query image, a departure
from the more straightforward foreground-background dichotomy. Multi-class FSS
typically adopts an $N$-way $K$-shot framework, entailing $N$ classes within
the support set and $K$ instances per class. This approach, geared towards
segmenting multiple classes, invariably relies on a prototype-based model
architecture to encapsulate and convey class-specific information throughout
the network. For example, MFNet [42] leverages Masked Global Pooling alongside
attention mechanisms to generate class prototypes that integrate with query
image features for segmentation.

It is crucial to recognize that the efficacy of multi-way methodologies is
contingent upon predefined $N$ and $K$ values, necessitating their
specification prior to training—a constraint that also influences many single-
way models in terms of the number of examples per segmented concept.

###  2.2 Promptable Segmentation

Integrating prompt-based learning, inspired by its success in natural language
processing, into computer vision introduced a novel paradigm. In particular, a
significant breakthrough came from the Segment Anything model (SAM),
introduced by Kirillov et al. [9]. This innovation harnesses an interactive,
prompt-guided vision framework, representing a pivotal moment reminiscent of
the GPT series but in computer vision. SAM has been trained on a vast dataset
comprising over 1 billion masks extracted from 11 million images, employing a
“promptable” segmentation task. This innovative approach equips SAM with the
unique capability for robust zero-shot generalization, enabling it to excel in
a wide range of tasks such as edge detection and instance segmentation.

In recent times, the research community has been prolific in extending the
capabilities of SAM and exploring its potential across diverse domains. These
endeavors have pushed the boundaries of what SAM can accomplish and have
successfully applied it to a wide range of tasks, including medical image
analysis [20], image inpainting [37], and image editing [36]. Recent
innovations like SegGPT [34] and SEEM [44] have extended the use of prompts to
guide predictions, exploring the integration of textual and audio prompts to
enhance model performance across diverse tasks.

However, despite these advancements, there remains a gap in applying these
models to FSS, especially in classifying objects into previously unseen
classes based on a limited number of annotated examples. While models like
SegGPT show promise in FSS scenarios, their generalization capabilities are
often attributed to the vast datasets used for training rather than an
inherent ability to leverage support set information effectively. Furthermore,
apart from an early exploration by Rakelly et al. [24], research in FSS has
scarcely investigated the use of diverse annotations within support sets—such
as points or bounding boxes—as alternatives to traditional mask prompts,
leaving a substantial area of FSS largely unexplored.

###  2.3 Our Contribution

Our work seeks to bridge the gaps in the literature using various visual
prompts to guide the FSS task. Drawing from cutting-edge prompt-based
segmentation techniques, LA innovatively employs an array of prompt types
within the FSS framework. Our novel end-to-end neural network architecture is
also designed for multi-class FSS, enhancing the versatility of forming
support sets, crucial for effective FSS. Moving beyond the conventional focus
on binary FSS, which mainly leverages affinity-based methods, our work
proposes a prototype-based architecture. This design choice aligns with the
efficiency and coherence demands of multi-way scenarios and addresses the
limitations imposed by the significant memory requirements of affinity-based
approaches in $N$-way $K$-shot setups. Taking cues from successful $1$-way
$K$-shot models (e.g., [21, 31, 38, 27]), which adapt to generic $K$-shot
scenarios post-training, our method expands this approach to accommodate
multi-class contexts. We train our model once, enabling it to seamlessly adapt
to various FSS benchmarks, including $1$-way $1$-shot and $N$-way $K$-shot
configurations, without imposing any constraints on the number of examples
($K$) per class. This approach affords unparalleled flexibility, allowing the
model to interpret support sets comprising any number of images and
annotations, irrespective of class or prompt type. Such enhancements
inherently reduce computational demands and significantly boost the model’s
versatility and robustness in diverse FSS scenarios.

##  3 Label Anything

We aim to develop a model capable of labeling any image across any class using
a minimal set of prompts for each class. This approach does not prescribe a
predefined number of prompts per class; instead, it mandates that the support
set comprehensively covers all required classes.

###  3.1 Problem Formulation

We introduce the task of multi-prompt, multi-way, few-shot semantic
segmentation. This involves classifying the pixels of a query image $I_{q}$
into one of $N+1$ classes, with $N$ denoting the number of foreground
categories to be segmented. The $N$ foreground concepts are learned using only
a few annotated images. To achieve this goal, we employ “episodic training”
[32], which consists of training episodes involving a query image and a
distinct support set.

Formally, our dataset is partitioned into two subsets: a training set
$\mathcal{D}_{train}$ and a test set $\mathcal{D}_{test}$. Crucially, we
define two distinct sets of segmentation categories, $\mathcal{C}_{seen}$ for
training and $\mathcal{C}_{unseen}$ for testing, ensuring
$\mathcal{C}_{seen}\cap\mathcal{C}_{unseen}=\emptyset$. Annotations in
$\mathcal{D}_{train}$ utilize classes from $\mathcal{C}_{seen}$, whereas
$\mathcal{D}_{test}$ is evaluated on classes from $\mathcal{C}_{unseen}$,
representing concepts not encountered during training to assess the model’s
generalization capabilities in a few-shot setting.

Each training episode presents the model with a query and its corresponding
ground truth $(I_{q},M_{q})$, where $I_{q}$ is an RGB image of dimensions
$3\times H\times W$ and $M_{q}$ is a segmentation mask associating pixels of
$I_{q}$ with one of $N+1$ classes. The model receives a support set
$\mathcal{S}=\\{(I_{i},A_{i})\\}_{i=1}^{L}$ comprising $L$ image-prompt
annotation pairs to assist in classifying the query pixels. Each pair
$(I_{i},A_{i})$ is called a “shot”, with $L$ usually being a small number. For
instance, in multi-class FSS tasks, $L$ is often determined as $N\times K$,
with $K$ indicating the number of examples, or shots, per class necessary for
the model to learn about class $n$. Standard testing configurations on
benchmark datasets typically involve setting $N$ and $K$ to values from the
set $\\{1,2,5\\}$.

Our approach significantly extends the flexibility of FSS by allowing for
multiple types of prompts—masks, points, and bounding boxes—within the support
set. Unlike previous FSS methodologies, LA does not require the support set
size $L$ to match $N\times K$ strictly, enabling variable numbers of shots per
class. Additionally, LA permits a single image to serve as a support example
for multiple classes without repetition, linking different class prompts to
the same image, thus maximizing the exploitation of inter-class relationships
within annotations.

In more detail, we have $A_{i}=\\{(M_{i,n},P_{i,n},B_{i,n})\\}_{n=1}^{N}$ for
each support set image, where $M_{i,n}$ is a binary mask for class $n$ in
image $i$, $P_{i,n}$ represents point coordinates, and $B_{i,n}$ denotes
bounding boxes. The prompt type for class $n$ in image $i$ is randomly
selected during training, with padding applied as necessary for absent classes
or prompt types. Evaluation on $\mathcal{D}_{test}$ involves support set
annotations from $\mathcal{C}_{unseen}$, testing the model’s adaptability to
new classes.

Figure 2: Class prototype generation process. LA begins with extracting image
features from a single support set example using ViT. This example might
include multiple prompts related to different classes and types. A CNN-based
Mask Encoder processes dense prompts (masks), while point coordinates and
bounding boxes are transformed into embeddings via the Point Encoder, with
bounding boxes treated as pairs of point embeddings. These processed prompts
are then integrated with the image features. The Class Token Pool enriches
dense and sparse embeddings with class information to infuse class-specific
semantics. Subsequent two-way cross-attention between dense and sparse
features yields class prototypes for each example.

###  3.2 Model Architecture

Label Anything processes a query image $I_{q}$ for segmentation alongside a
support set $\mathcal{S}$, utilizing three core components in its architecture
(refer to Fig. 1): the Image Encoder, the Prompt Encoder, and the Mask
Decoder.

The Image Encoder leverages a pre-trained Vision Transformer, specifically
ViT-B/16 [5] to process the query and support images. These images are
transformed into feature maps $F_{q}$ for the query and $F_{s}$ for the
support set, with dimension $D_{ViT}\times H_{d}\times W_{d}$, where
$H_{d}=H/x_{p}$ and $W_{d}=W/x_{p}$ and $x_{p}$ is the ViT patch size.
Attached to ViT is a convolutional neck that processes the image features to
reduce their dimensionality from $D_{ViT}=768$ to $D=512$.

The Prompt Encoder is a key component of the LA architecture, designed to
embody different types of prompts—masks, boxes, and points—from support images
and effectively encapsulate this information. In its initial stage, the Prompt
Encoder processes these inputs to generate embeddings: masks are encoded using
a CNN, resulting in dense embeddings ($P_{d}$), while points and boxes are
encoded through a Point Encoder, which uses Positional Embeddings to produce
sparse embeddings for points and boxes ($P_{s}$). Three different learnable
embeddings are used, each corresponding to the three types of coordinate
inputs (point prompt, box top-left corner, box bottom-right corner), which are
then summed. To enable attention on various objects as dictated by sparse
prompt features, a self-attention layer complemented by a residual connection
is employed. The layer operates on an input dimensionality of $BL\times
NM\times D$, signifying batch size ($B$), number of examples ($L$), classes
($N$), prompts per class ($M$), and feature size ($D$), thereby facilitating
independent example processing and attention to each prompt feature across
both similar and different classes. We adopt the Token Pool strategy from
Zhang et al. [41] to incorporate class-specific information using a trainable
random Gaussian matrix sized $R\times D$, where $R$ should be greater than the
number of classes. At each step, a row is sampled from this matrix for each
class and added to the sparse and dense features. Subsequently, dense
embeddings are combined with the image features to highlight the featured
object within the feature map. Merging sparse and dense embeddings in a single
feature map requires a Two-Way Transformer, similar to the mechanism in SAM
[9]. This Transformer alternates attention between dense to sparse features
and vice versa, according to:

$\displaystyle P_{s}$
$\displaystyle=MultiHeadAttention(P_{s},P_{d}+PE,P_{d}+PE),$ (1)
$\displaystyle P_{d}$
$\displaystyle=MultiHeadAttention(P_{d}+PE,P_{s},P_{s}),$ (2)

where $PE$ denotes the positional encoding adjusted for dimensions $(H,W)$.
This process ensures a comprehensive integration of prompt-based information,
highlighting the relevance of each object within the feature map.

After processing sparse features, these features are eliminated, forcing the
model to encapsulate all pertinent details within the dense features. A Global
Average Pooling layer then condenses the spatial dimensions, generating a set
of class-example embeddings $E=\\{E_{ni}\in\mathbb{R}^{D}|\forall
n\in\mathcal{C},\forall i\in\mathcal{S}\\}$. To facilitate the interchange of
information across examples and classes, we introduce the Class Example Mixer
block, which enables communication between different examples through a
Transformer layer. The attention layer in it enables the embeddings of the
same class $n$ over different examples to be attended:

$\displaystyle E_{n}$ $\displaystyle=E_{n_{1}},E_{n_{2}},\dots,E_{n_{L}}$ (3)
$\displaystyle E_{n}$ $\displaystyle=MultiHeadAttention(E_{n},E_{n},E_{n})$
(4)

Embeddings are averaged across examples to yield a singular prototype per
class, compacting essential class information. Figure 2 shows a detailed view
of the class prototype generation for a single support set example.

The Mask Decoder uses the class prototypes $E_{1},\dots,E_{|\mathcal{C}|}$ (or
class-example prototypes at test time), and the query image features $F_{q}$
to generate a segmentation mask for each class. This process begins with the
Two-Way Transformer, which serves a role different from its Prompt Encoder
function. Here, it aims to facilitate pattern matching between class
prototypes and pixel locations within the query image by reciprocally
attending to them, thus transferring relevant information for segmentation.
Subsequently, the query features are upscaled to a resolution of $(H/4,W/4)$,
and their dimensionality is reduced to $D/8$ through a sequence of transposed
convolutional layers. Concurrently, the class prototypes are dimensionally
aligned to match the query features via an MLP, ensuring uniformity in
representation. Following this, the modified query features $F_{q}$ undergo
further refinement through three $3\times 3$ spatial convolutional layers
aimed at producing segmentation outputs that are not only smoother but also
spatially coherent. The segmentation logits for each pixel and class are
determined through a dot product operation, succinctly captured by the
formula:

$\text{Logit}_{xy,n}=F_{q_{xy}}\cdot E_{n},$ (5)

thereby computing the final segmentation masks.

###  3.3 Training Procedure

One of the primary objectives of this work is to enable our model to be
trained once for application across arbitrary combinations of $N$ and $K$
during testing and across various prompt inputs. To accomplish this, we
introduce an innovative episodic training approach that optimizes resource
utilization while accommodating diverse $N$, $K$, and prompt-type
configurations.

The procedure starts by defining potential batch configurations as $(B,N,K)$,
where $B$ represents the batch size, indicating the number of learning
episodes per batch. The selection of $B$ is tailored to the $N$ and $K$
parameters to ensure optimal GPU resource management using powers of 2, with
$B\times N\times K+B$ representing the aggregate image count per batch,
including both support and query images. Each training batch is then
constructed by randomly selecting a $(B,N,K)$ tuple.

For every episode within a batch, $N$ classes are randomly chosen from the set
of seen classes, $\mathcal{C}_{seen}$. A coin flip probabilistic decision
determines whether the query image should contain any or all of the $N$
selected classes, adjusting $N$ if a comprehensive query image is unavailable.
This variability fosters episodes where the query image encompasses objects
from all sampled classes and scenarios, necessitating the exclusion of certain
support set classes during segmentation, thereby addressing a recognized bias
in FSS research [8].

Subsequently, $K$ support examples per class are sampled. We identify all
annotations corresponding to the selected $N$ classes within each chosen
image, allowing multiple annotations per image. This flexibility ensures our
model’s readiness for any $L$ image-annotation pair configuration, where $L$
may deviate from the strict $N\times K$ formula.

Additionally, our training episodes are enriched by incorporating a random
selection of prompt types for each instance annotation within the support set,
enhancing the diversity of learning stimuli. For instances requiring point
annotations, points are sampled proportionately to the instance’s mask area,
tailoring the learning process to the specific needs of each episode and
prompt type.

##  4 Experimental Evaluation

We evaluated LA across two distinct FSS tasks: $1$-way $1$-shot segmentation
and $N$-way $K$-shot segmentation. Performance is reported as mean
Intersection-over-Union (mIoU) across test classes.

###  4.1 Dataset and Validation

Our assessment was based on the COCO-$20^{i}$ dataset [22, 33], a predominant
FSS benchmark sourced from the MS COCO dataset [12], to evaluate our model
across both $1$-way $1$-shot and $N$-way $K$-shot segmentation tasks.
COCO-$20^{i}$ amplifies the challenge by dividing its 80 classes into four
distinct folds, each encompassing 20 classes. Employing cross-validation, we
rotated classes within each fold as $\mathcal{C}_{unseen}$, contrasting with
$\mathcal{C}_{seen}$ for the remaining classes. This setup aligned
$\mathcal{D}_{train}$ and $\mathcal{D}_{test}$ with the original dataset’s
arrangement, ensuring our model faced novel images and classes during testing.
Our model was fine-tuned in a single iteration per cross-validation phase,
highlighting our training strategy’s effectiveness.

To ensure a rigorous comparison with established benchmarks, especially in
binary FSS contexts [21, 22], we adopted a consistent evaluation methodology,
conducting tests over $1000$ randomly chosen episodes, utilizing $5$ distinct
random seeds for each and averaging the outcomes. This procedure was
replicated for $N$-way $K$-shot evaluations, aligning with methodologies from
prominent studies [42, 33, 14]. Unless otherwise stated, we limited our prompt
type to masks during tests for more accurate benchmarking, facilitating direct
comparisons across different approaches and configurations.

###  4.2 Setting

We adopted a uniform input resolution of $1024\times 1024$ for query and
support images, resizing images and applying padding to maintain the aspect
ratio. Following the methodology proposed by Zhang et al. [41], we used the
ViT-B/16 model [5], augmented with Segment Anything [9] pretraining and with
its weights kept frozen, serving as our backbone. This configuration yielded a
singular feature map per image. Class-specific masks were supplied at a
refined resolution of 256 $\times$ 256, while points and bounding boxes were
integrated into the model using normalized coordinates.

Moreover, for a fairer comparison with existing methods, we also employed the
same ViT-B/16 architecture with the Masked Autoencoder (MAE) pre-training
strategy [2]. The input size was $480\times 480$, with the output feature map
of dimension $768\times 30\times 30$. This pretraining strategy is preferred
against the classical ImageNet classifiation pre-training as it retains the
spatial features useful for segmentation.

Our training strategy optimized the focal loss [11], which addresses class
imbalance by focusing on difficult-to-classify examples computed for a single
training instance as:

$\mathcal{L}=\frac{1}{N+1}\sum_{n}^{N+1}{\left[w_{n}\cdot(1-e^{-l_{ce}(\hat{y}_{n},y_{n})})^{\gamma}\cdot
l_{ce}(\hat{y}_{n},y_{n})\right]},$ (6)

where $N+1$ is the number of classes, $w_{n}$ are class-specific weights,
$l_{ce}$ the cross-entropy loss, and $\gamma$ the focusing parameter adjusting
the emphasis on hard examples.

We used the AdamW optimizer [19] ($\beta_{1}=0.9$, $\beta_{2}=0.999$),
complemented by a linear learning rate warmup [7] for 1000 iterations,
followed by a step-wise cosine learning rate decay schedule [18]. The initial
learning rate, after warmup, was $1\mathrm{e}{-5}$. Our training followed the
configurations outlined in Sec. 3.3, adopting batch size-number of classes-
number of shots configurations such as $(4,1,4)$, $(2,4,2)$, $(8,1,2)$,
$(4,2,2)$, $(4,4,1)$, and $(16,1,1)$, where the batch size refers to a single
GPU. The model underwent training for 50 epochs, with the iteration count per
epoch equivalent to the size of $\mathcal{D}_{train}$. A cap was placed on
point samples at a maximum of $10$ per instance to maintain efficiency. The
architecture totals 27 million learnable parameters, reflecting the model’s
complexity and capacity for handling diverse FSS scenarios.

Training leveraged the Leonardo cluster’s resources, including 512GB of RAM
and four NVIDIA A100-64GB GPUs, to support the computational demands, ensuring
high-performance processing capabilities.

Method |  $1$-way $1$-shot |  $1$-way $5$-shot
---|---|---
fold-0 | fold-1 | fold-2 | fold-3 | mean | fold-0 | fold-1 | fold-2 | fold-3 | mean
DCAMA [27] | 49.5 | 52.7 | 52.8 | 48.7 | 50.9 | 55.4 | 60.3 | 59.9 | 57.5 | 58.3
FPTrans [40] | 39.7 | 44.1 | 44.4 | 39.7 | 42.0 | 49.9 | 56.5 | 55.4 | 53.2 | 53.8
LA (MAE) | 35.2 | 44.0 | 39.4 | 40.4 | 39.7 | 38.5 | 44.8 | 39.5 | 42.3 | 41.2
LA (SAM) | 39.2 | 46.6 | 44.1 | 42.7 | 43.1 | 42.4 | 49.3 | 45.9 | 43.4 | 45.1

Table 1: Results for $1$-way $K$-shot segmentation, benchmarking against
leading transformer-based FSS methodologies. DCAMA [27] represents the peak in
affinity learning-based segmentation, leveraging the hierarchical Vision
Transformer Swin-B [16] as its foundation. Conversely, FPTrans [41]
exemplifies top-tier performance in prototype-based segmentation, employing
the ViT-B/16 [5] as its backbone.

###  4.3 Results

In Table 1, we showcase the performance of LA on the COCO-$20^{i}$ benchmark,
juxtaposing it with leading binary FSS methods. Our model achieved state-of-
the-art results among prototype-based approaches, notably surpassing FPTrans
[41] in the 1-shot setting. However, DCAMA [27], an affinity-based method,
exceeds our performance by comparing each pixel in the query image against
every pixel in the support images, leveraging features from multiple backbone
levels. Given its detailed pixel-wise comparison strategy, DCAMA’s superior
performance was anticipated. Despite this, affinity-based methods face
scalability challenges, particularly in $N$-way $K$-shot scenarios, as the
number of necessary comparisons increases quadratically with the image count
and, by extension, with the class count. This limitation underscores the
practical advantages of prototype-based models like LA for complex
segmentation tasks involving multiple classes and shots.

Indeed, Table 2 presents the performance of LA and several other state-of-the-
art methods on the $2$-way $K$-shot segmentation task, as evaluated on the
COCO-$20^{i}$ benchmark. LA demonstrates superior performance across all four
folds, excelling in the more challenging $2$-way $1$-shot and $2$-way $5$-shot
scenarios, achieving the highest mIoU scores. Notably, LA surpassed MFNet
[42], previously regarded as the leading method, especially in folds 1, 2, and
3, underscoring its superior ability to handle multi-class segmentation tasks
with limited examples. Other methods show competitive but lower performance
compared to LA, highlighting the effectiveness of our approach in leveraging
diverse visual prompts and an end-to-end training strategy.

Furthermore, to validate the performances of LA in complex FSS settings, we
tested our model capabilities also on an increasing $N$ setting. Table 3 shows
the results of LA and DCAMA on the $N$-way $1$-shot segmentation task. Except
for 1-way, LA outperforms DCAMA in all other $N$-way settings, showcasing its
robustness and adaptability with a growing number of classes.

Figure 3 presents qualitative results from our $2$-way $1$-shot segmentation
experiment on the COCO-$20^{i}$ dataset. This visual representation highlights
the efficacy of our model in accurately interpreting and responding to diverse
prompts, showcasing its robust segmentation capabilities in complex scenarios.

Method |  $2$-way $1$-shot |  $5$-way $1$-shot
---|---|---
fold-0 | fold-1 | fold-2 | fold-3 | mean | fold-0 | fold-1 | fold-2 | fold-3 | mean
PANet [33] | 25.7 | 16.1 | 16.2 | 13.8 | 18.0 | 22.1 | 16.0 | 14.1 | 11.9 | 16.0
PPNet(w/o U) [14] | 29.0 | 19.4 | 16.5 | 14.2 | 19.8 | 24.3 | 16.8 | 14.3 | 12.8 | 17.1
PPNet [14] | 29.8 | 19.7 | 17.0 | 15.1 | 20.4 | 25.6 | 17.3 | 15.5 | 13.4 | 18.0
DENet [13] | 30.3 | 20.5 | 16.7 | 15.3 | 20.7 | 27.0 | 17.9 | 14.6 | 14.0 | 18.4
MFNet [42] | 35.3 | 24.0 | 18.4 | 18.7 | 24.1 | 29.7 | 21.2 | 15.4 | 17.1 | 20.9
LA (MAE) | 30.6 | 33.6 | 32.2 | 31.2 | 31.9 | 21.1 | 25.1 | 26.4 | 27.5 | 25.1
LA (SAM) | 31.0 | 38.2 | 34.7 | 34.5 | 34.6 | 24.2 | 27.5 | 29.6 | 29.7 | 27.7

Table 2: Performance on $2$-way / $5$-way $1$-shot segmentation benchmarked with contemporaneous methods as re-evaluated in the MFNet study [42]. Comparisons are made using mIoU percentages, specifically focusing on foreground classes. Misclassifications within these classes are considered under false positive and false negative categories for the respective classes; a measurement approach denoted as mIoU* in the MFNet analysis. Method | 1-way | 2-way | 3-way | 4-way | 5-way | 10-way | 15-way | 20-way
---|---|---|---|---|---|---|---|---
DCAMA | 50.9 | 31.7 | 24.2 | 21.1 | 16.7 | 10.8 | 6.5 | 4.7
LA (MAE) | 39.7 | 31.9 | 28.6 | 26.5 | 25.1 | 20.7 | 15.0 | 10.8
LA (SAM) | 43.1 | 34.6 | 31.7 | 29.6 | 27.7 | 23.6 | 16.9 | 13.7
Table 3: Average mIoU results for $N$-way 1-shot on COCO-$20^{i}$.

Figure 3: Visual representation of $2$-way $1$-shot segmentation on
COCO-$20^{i}$. Sequence from left to right includes the first prompt, the
second prompt, the ground truth mask, and our model’s prediction.

###  4.4 Ablation Study

Our comprehensive ablation studies on the validation fold of COCO-$20^{i}$
dissected the contribution of each component within our model to its overall
performance. We systematically omitted elements deemed non-
critical—specifically, the Spatial Convolutions in the Mask Decoder, the Class
Token Pool, and the Class Example Mixer from the Prompt Encoder—and observed
the resultant impact on model efficacy. As delineated in Table 4, excluding
the Spatial Convolutions from the Mask Decoder markedly diminished model
performance. This component is pivotal for blending spatial details before
classification, significantly enhancing segmentation quality. Meanwhile, the
Class Token Pool and Class Example Mixer, though their removal slightly
affected performance, were proven indispensable for optimizing results. These
findings highlight the integral role of each component in achieving superior
segmentation outcomes, underscoring the holistic synergy that drive our
model’s performance.

Our analysis also assessed the impact of different prompt types on our model’s
effectiveness, exploring how each type influences performance. Table 5 reveals
that mask prompts led to the highest performance, offering the most detailed
guidance for segmentation tasks. Nonetheless, the model demonstrated robust
adaptability, yielding solid results even with alternative prompt types,
albeit with a marginal reduction compared to masks alone. This underscores the
model’s versatile capability to effectively manage and interpret various
prompt types.

Configuration | 1-way 1-shot | 2-way 1-shot
---|---|---
Complete model | 39.2 | 30.9
w/o Spatial Convolutions | 34.2 | 27.6
w/o Class Example Mixer | 38.7 | 30.9
w/o Token Pool | 38.3 | 30.1
Table 4: Ablation analysis on the COCO-$20^{i}$’s fold-0 evaluating the impact of different components of the Label Anything architecture. Prompt Type | 1-way 1-shot | 1-way 5-shot | 2-way 1-shot | 2-way 5-shot
---|---|---|---|---
Only masks | 43.5 | 45.1 | 34.6 | 35.9
Boxes and points | 37.1 | 40.8 | 29.2 | 30.7
Only boxes | 38.3 | 41.2 | 30.6 | 31.5
Only points | 36.5 | 39.9 | 29.0 | 29.4
Table 5: Ablation analysis evaluating the impact of different prompt types on
model performance.

##  5 Conclusion

In this study, we presented Label Anything, an innovative framework for few-
shot semantic segmentation that markedly advances the adaptability and
efficiency of segmentation models. Our rigorous testing on the COCO-$20^{i}$
benchmark has showcased the model’s capacity to reach competitive and, in some
instances, state-of-the-art performance, alongside its ability to generalize
across diverse segmentation scenarios. Including a wide range of visual
prompts—points, bounding boxes, and masks—empowers Label Anything to interpret
and leverage multiple forms of input guidance effectively. The results
highlight the efficacy of prototype-based approaches in addressing the
challenges of multi-class FSS, presenting a compelling alternative to
affinity-based methods which, despite their accuracy, struggle with
scalability as the dataset expands.

We envision a broad spectrum of future research directions to advance our
method. An intriguing avenue involves the integration of textual prompts,
leveraging the interplay between textual and visual information to enrich the
segmentation process. Concurrently, the proposal of an optimized example
candidate selection mechanism, potentially through clustering approaches,
promises to streamline the identification of example candidates within large-
scale datasets, enhancing the efficiency of support set construction.

####  5.0.1 Acknowledgement

We acknowledge the CINECA award under the ISCRA initiative, which gave us
access to computing resources and support.

## References

  * [1] Chen, H., Dong, Y., Lu, Z., Yu, Y., Han, J.: Pixel Matching Network for Cross-Domain Few-Shot Segmentation. In: Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. pp. 978–987 (2024)
  * [2] Cheng, B., Misra, I., Schwing, A.G., Kirillov, A., Girdhar, R.: Masked-Attention Mask Transformer for Universal Image Segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 1290–1299 (2022)
  * [3] Çiçek, Ö., Abdulkadir, A., Lienkamp, S.S., Brox, T., Ronneberger, O.: 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19. pp. 424–432. Springer (2016)
  * [4] Dong, N., Xing, E.P.: Few-shot semantic segmentation with prototype learning. In: BMVC (2018)
  * [5] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.: An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 (2020)
  * [6] Finn, C., Abbeel, P., Levine, S.: Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks (Jul 2017), arXiv:1703.03400 [cs]
  * [7] Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., He, K.: Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 (2017)
  * [8] Kang, D., Cho, M.: Integrative few-shot learning for classification and segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 9979–9990 (2022)
  * [9] Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A.C., Lo, W.Y., et al.: Segment Anything. arXiv preprint arXiv:2304.02643 (2023)
  * [10] Li, G., Jampani, V., Sevilla-Lara, L., Sun, D., Kim, J., Kim, J.: Adaptive prototype learning and allocation for few-shot segmentation. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 8334–8343 (2021)
  * [11] Lin, T.Y., Goyal, P., Girshick, R., He, K., Dollár, P.: Focal loss for dense object detection. In: Proceedings of the IEEE international conference on computer vision. pp. 2980–2988 (2017)
  * [12] Lin, T.Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., Zitnick, C.L.: Microsoft coco: Common objects in context. In: Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13. pp. 740–755. Springer (2014)
  * [13] Liu, L., Cao, J., Liu, M., Guo, Y., Chen, Q., Tan, M.: Dynamic extension nets for few-shot semantic segmentation. In: Proceedings of the 28th ACM international conference on multimedia. pp. 1441–1449 (2020)
  * [14] Liu, Y., Zhang, X., Zhang, S., He, X.: Part-aware prototype network for few-shot semantic segmentation. In: Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part IX 16. pp. 142–158. Springer (2020)
  * [15] Liu, Y., Liu, N., Yao, X., Han, J.: Intermediate prototype mining transformer for few-shot semantic segmentation. Advances in Neural Information Processing Systems 35, 38020–38031 (2022)
  * [16] Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin transformer: Hierarchical vision transformer using shifted windows. In: Proceedings of the IEEE/CVF international conference on computer vision. pp. 10012–10022 (2021)
  * [17] Long, J., Shelhamer, E., Darrell, T.: Fully Convolutional Networks for Semantic Segmentation. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 3431–3440 (2015)
  * [18] Loshchilov, I., Hutter, F.: Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 (2016)
  * [19] Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 (2017)
  * [20] Mazurowski, M.A., Dong, H., Gu, H., Yang, J., Konz, N., Zhang, Y.: Segment Anything Model for Medical Image Analysis: An Experimental Study. Medical Image Analysis 89, 102918 (2023)
  * [21] Min, J., Kang, D., Cho, M.: Hypercorrelation squeeze for few-shot segmentation. In: Proceedings of the IEEE/CVF international conference on computer vision. pp. 6941–6952 (2021)
  * [22] Nguyen, K., Todorovic, S.: Feature weighting and boosting for few-shot segmentation. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 622–631 (2019)
  * [23] Oktay, O., Schlemper, J., Folgoc, L.L., Lee, M., Heinrich, M., Misawa, K., Mori, K., McDonagh, S., Hammerla, N.Y., Kainz, B., et al.: Attention U-Net: Learning Where to Look for the Pancreas. arXiv preprint arXiv:1804.03999 (2018)
  * [24] Rakelly, K., Shelhamer, E., Darrell, T., Efros, A.A., Levine, S.: Conditional Networks for Few-Shot Semantic Segmentation. In: International Conference on Learning Representations (2018)
  * [25] Ronneberger, O., Fischer, P., Brox, T.: U-Net: Convolutional Networks for Biomedical Image Segmentation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. pp. 234–241. Springer (2015)
  * [26] Shaban, A., Bansal, S., Liu, Z., Essa, I., Boots, B.: One-Shot Learning for Semantic Segmentation. British Machine Vision Conference (BMVC) (2017)
  * [27] Shi, X., Wei, D., Zhang, Y., Lu, D., Ning, M., Chen, J., Ma, K., Zheng, Y.: Dense cross-query-and-support attention weighted mask aggregation for few-shot segmentation. In: European Conference on Computer Vision. pp. 151–168. Springer (2022)
  * [28] Snell, J., Swersky, K., Zemel, R.S.: Prototypical Networks for Few-shot Learning (Jun 2017), arXiv:1703.05175 [cs, stat]
  * [29] Sung, F., Yang, Y., Zhang, L., Xiang, T., Torr, P.H.S., Hospedales, T.M.: Learning to Compare: Relation Network for Few-Shot Learning (Mar 2018), arXiv:1711.06025 [cs]
  * [30] Tian, P., Wu, Z., Qi, L., Wang, L., Shi, Y., Gao, Y.: Differentiable meta-learning model for few-shot semantic segmentation. In: Proceedings of the AAAI Conference on Artificial Intelligence. pp. 12087–12094 (2020)
  * [31] Tian, Z., Zhao, H., Shu, M., Yang, Z., Li, R., Jia, J.: Prior guided feature enrichment network for few-shot segmentation. IEEE transactions on pattern analysis and machine intelligence 44(2), 1050–1065 (2020)
  * [32] Vinyals, O., Blundell, C., Lillicrap, T., Kavukcuoglu, K., Wierstra, D.: Matching Networks for One Shot Learning (Dec 2017), arXiv:1606.04080 [cs, stat]
  * [33] Wang, K., Liew, J.H., Zou, Y., Zhou, D., Feng, J.: Panet: Few-shot image semantic segmentation with prototype alignment. In: proceedings of the IEEE/CVF international conference on computer vision. pp. 9197–9206 (2019)
  * [34] Wang, X., Zhang, X., Cao, Y., Wang, W., Shen, C., Huang, T.: SegGPT: Segmenting Everything In Context. arXiv e-prints pp. arXiv–2304 (2023)
  * [35] Wang, Y., Sun, R., Zhang, T.: Rethinking the Correlation in Few-Shot Segmentation: A Buoys View. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 7183–7192 (2023)
  * [36] Xie, D., Wang, R., Ma, J., Chen, C., Lu, H., Yang, D., Shi, F., Lin, X.: Edit Everything: A Text-Guided Generative System for Images Editing. arXiv preprint arXiv:2304.14006 (2023)
  * [37] Yu, T., Feng, R., Feng, R., Liu, J., Jin, X., Zeng, W., Chen, Z.: Inpaint Anything: Segment Anything Meets Image Inpainting. arXiv preprint arXiv:2304.06790 (2023)
  * [38] Zhang, B., Xiao, J., Qin, T.: Self-guided and cross-guided learning for few-shot segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 8312–8321 (2021)
  * [39] Zhang, C., Lin, G., Liu, F., Yao, R., Shen, C.: Canet: Class-agnostic segmentation networks with iterative refinement and attentive few-shot learning. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 5217–5226 (2019)
  * [40] Zhang, G., Kang, G., Yang, Y., Wei, Y.: Few-shot segmentation via cycle-consistent transformer. Advances in Neural Information Processing Systems 34, 21984–21996 (2021)
  * [41] Zhang, J.W., Sun, Y., Yang, Y., Chen, W.: Feature-proxy transformer for few-shot segmentation. Advances in Neural Information Processing Systems 35, 6575–6588 (2022)
  * [42] Zhang, M., Shi, M., Li, L.: MFNet: Multiclass Few-Shot Segmentation Network With Pixel-Wise Metric Learning. IEEE Transactions on Circuits and Systems for Video Technology 32(12), 8586–8598 (2022)
  * [43] Zhang, X., Wei, Y., Yang, Y., Huang, T.S.: Sg-one: Similarity guidance network for one-shot semantic segmentation. IEEE transactions on cybernetics 50(9), 3855–3865 (2020)
  * [44] Zou, X., Yang, J., Zhang, H., Li, F., Li, L., Wang, J., Wang, L., Gao, J., Lee, Y.J.: Segment everything everywhere all at once. Advances in Neural Information Processing Systems 36 (2024)