Prototype Correlation Matching and Class-Relation Reasoning for Few-Shot Medical Image Segmentation

TL;DR

This work tackles the challenge of generalizing few-shot medical image segmentation to unseen classes despite large intra-class variations. It introduces PMCR, a framework that integrates Prototype Correlation Matching (PCM) to mine representative prototypes and perform prototype-level optimal-transport matching, with Class-Relation Reasoning (CRR) to propagate inter-class relations via a superpixel-based relation graph and kernel attention. Key contributions include SVD-derived prototypes, OT-based matching, a memory-augmented superpixel graph, and a convolution-kernel attention scheme that refines query encoding using base-novel class relations; these yield substantial improvements on Cardiac-MRI, Abdominal-MRI, Abdominal-CT, and Prostate-MRI datasets. The approach demonstrates strong generalization to unseen classes while maintaining practical efficiency, highlighting its potential for clinical deployment and future extensions to federated, continual, and multi-modality settings.

Abstract

Few-shot medical image segmentation has achieved great progress in improving accuracy and efficiency of medical analysis in the biomedical imaging field. However, most existing methods cannot explore inter-class relations among base and novel medical classes to reason unseen novel classes. Moreover, the same kind of medical class has large intra-class variations brought by diverse appearances, shapes and scales, thus causing ambiguous visual characterization to degrade generalization performance of these existing methods on unseen novel classes. To address the above challenges, in this paper, we propose a \underline{\textbf{P}}rototype correlation \underline{\textbf{M}}atching and \underline{\textbf{C}}lass-relation \underline{\textbf{R}}easoning (i.e., \textbf{PMCR}) model. The proposed model can effectively mitigate false pixel correlation matches caused by large intra-class variations while reasoning inter-class relations among different medical classes. Specifically, in order to address false pixel correlation match brought by large intra-class variations, we propose a prototype correlation matching module to mine representative prototypes that can characterize diverse visual information of different appearances well. We aim to explore prototype-level rather than pixel-level correlation matching between support and query features via optimal transport algorithm to tackle false matches caused by intra-class variations. Meanwhile, in order to explore inter-class relations, we design a class-relation reasoning module to segment unseen novel medical objects via reasoning inter-class relations between base and novel classes. Such inter-class relations can be well propagated to semantic encoding of local query features to improve few-shot segmentation performance. Quantitative comparisons illustrates the large performance improvement of our model over other baseline methods.

Figures (9)

• Figure 1: Illustration of two major challenges in the few-shot medical image segmentation. (a) Inbuilt inter-class relations between different medical classes in the pure semantic space. (b) Large intra-class variations within the same medical class under different imaging protocols (e.g., the hepatocellular carcinoma (HCC) has various appearances 10.1093/bib/bbaa295).
• Figure 2: Pipeline of the proposed model. It includes a prototype correlation matching (PCM) module to address false pairwise pixel correlation matches brought by large intra-class variations within the same medical class; and a class-relation reasoning (CRR) module to explore inter-class relations between base and novel classes for semantic encoding of local query features.
• Figure 3: Exemplar segmentation results of the proposed model and some representative methods on Abdominal-MRI KAVUR2021101950 under the setting #1.
• Figure 4: Exemplar segmentation results of the proposed model and some representative methods on Abdominal-CT 10.1007/978-3-030-58526-6_45 under the setting #1.
• Figure 5: Analysis of hyper-parameters $\{\lambda_1, \lambda_2\}$ on Cardiac-MRI 8458220 (top), Abdominal-MRI KAVUR2021101950 (middle) and Abdominal-CT 10.1007/978-3-030-58526-6_45 (bottom) under the setting #1. We analyze $\lambda_1$ when $\lambda_2=1.0$ and investigate $\lambda_2$ when $\lambda_1=0.5$.