Improving 3D Few-Shot Segmentation with Inference-Time Pseudo-Labeling

TL;DR

This work addresses data scarcity in 3D medical image segmentation under $N$-way $K$-shot FSS. It proposes an inference-time three-stage strategy that starts with Stage 1 prototypes from the support set, then uses Stage 2 confidence-aware pseudo-labeling to harvest informative regions from the unlabeled query to augment the support prototypes, and finally applies Stage 3 to perform segmentation with the augmented prototype set. The key contribution is that inference-time augmentation of prototypes leverages unlabeled query data without additional supervision or retraining. Experiments on abdominal CT and MRI demonstrate consistent Dice improvements and validate the effectiveness of windowed pseudo-label transfer and mixed prototype sets.

Abstract

In recent years, few-shot segmentation (FSS) models have emerged as a promising approach in medical imaging analysis, offering remarkable adaptability to segment novel classes with limited annotated data. Existing approaches to few-shot segmentation have often overlooked the potential of the query itself, failing to fully utilize the valuable information it contains. However, treating the query as unlabeled data provides an opportunity to enhance prediction accuracy. Specifically in the domain of medical imaging, the volumetric structure of queries offers a considerable source of valuable information that can be used to improve the target slice segmentation. In this work, we present a novel strategy to efficiently leverage the intrinsic information of the query sample for final segmentation during inference. First, we use the support slices from a reference volume to generate an initial segmentation score for the query slices through a prototypical approach. Subsequently, we apply a confidence-aware pseudo-labeling procedure to transfer the most informative parts of query slices to the support set. The final prediction is performed based on the new expanded support set, enabling the prediction of a more accurate segmentation mask for the query volume. Extensive experiments show that the proposed method can effectively boost performance across diverse settings and datasets.

Figures (2)

• Figure 1: Overview of the proposed method. The workflow begins with the extraction of embeddings from both the support and query slices using a shared feature extractor (Steps 1 and 2). In Step 3, support prototypes are generated from the support embeddings and corresponding ground truth masks. Step 4 calculates pseudo-masks for the query by measuring the distance between the support prototypes and the query embeddings. Next, in Step 5, query prototypes are generated from the query embeddings and pseudo-masks, and these are combined with the support prototypes in Step 6. Finally, in Step 7, the segmentation of the query slice is performed by using the augmented prototypes along with the query embeddings. Blue arrows depict the process of prototype calculation from feature embeddings and corresponding labels (masks), while red arrows indicate label prediction based on feature maps and prototypes. Background pixels are represented in black, foreground pixels in white, and low-confidence pixels in gray.
• Figure 2: Qualitative comparison of our method with SSLALPNet ouyang2020self