Concentrate on Weakness: Mining Hard Prototypes for Few-Shot Medical Image Segmentation

TL;DR

This work addresses boundary blur in few-shot medical image segmentation (FSMIS) by introducing Concentrate on Weakness (CoW), a framework that mines boundary-critical hard prototypes. CoW comprises Support Self-Prediction (SSP) to identify weak boundary features, Hard Prototypes Generation (HPG) to create multiple challenging prototypes, and Multiple Similarity Maps Fusion (MSMF) to fuse diverse similarity maps for accurate query segmentation, complemented by a boundary loss. Through extensive experiments on Abd-MRI, Abd-CT, and CMR datasets, CoW achieves state-of-the-art Dice scores, demonstrating improved boundary fidelity and robustness with minimal annotations. The approach advances practical FSMIS by more comprehensively representing class distributions and sharpening segmentation boundaries in medical images.

Abstract

Few-Shot Medical Image Segmentation (FSMIS) has been widely used to train a model that can perform segmentation from only a few annotated images. However, most existing prototype-based FSMIS methods generate multiple prototypes from the support image solely by random sampling or local averaging, which can cause particularly severe boundary blurring due to the tendency for normal features accounting for the majority of features of a specific category. Consequently, we propose to focus more attention to those weaker features that are crucial for clear segmentation boundary. Specifically, we design a Support Self-Prediction (SSP) module to identify such weak features by comparing true support mask with one predicted by global support prototype. Then, a Hard Prototypes Generation (HPG) module is employed to generate multiple hard prototypes based on these weak features. Subsequently, a Multiple Similarity Maps Fusion (MSMF) module is devised to generate final segmenting mask in a dual-path fashion to mitigate the imbalance between foreground and background in medical images. Furthermore, we introduce a boundary loss to further constraint the edge of segmentation. Extensive experiments on three publicly available medical image datasets demonstrate that our method achieves state-of-the-art performance. Code is available at https://github.com/jcjiang99/CoW.

Figures (9)

• Figure 1: Comparison between previous methods and ours. (a) The previous methods generate multiple prototypes only by random sampling or local averaging, giving insufficient attention to weak boundary details. (b) Our proposed CoW instead concentrates more on these weak features of a specific category, thus acquiring a more comprehensive representation of the category distribution and achieving clearer segmentation boundary.
• Figure 2: Illustration of CoW. We employ a shared feature encoder to learn deep features for both support and query images. In each episode, we firstly conduct SSP with support prototype $\mathbf{p}$ obtained via MAP. Next we compare this prediction with the ground truth $M_s$ to acquire the instruction for hard and normal feature points. Then we perform random sampling and reconstruct support feature map, which is used for generating multiple hard and normal prototypes for foreground and background, respectively. At last, these prototypes are applied to generate multiple prediction similarity maps that are fused by a lightweight decoder to get final segmentation mask.
• Figure 3: Comparison of qualitative results between our method and other on Abd-MRI.
• Figure 4: t-SNE visualization for generated normal and hard prototypes. 'Fg' implies foreground and 'Bg' implies background.
• Figure 5: Comparison of qualitative results between our method and other on Abd-CT.