Table of Contents
Fetching ...

Correlation Weighted Prototype-based Self-Supervised One-Shot Segmentation of Medical Images

Siladittya Manna, Saumik Bhattacharya, Umapada Pal

TL;DR

The paper tackles the challenge of medical image segmentation under limited annotations by proposing a self-supervised, one-shot framework that generates a dynamic prototype for every query pixel through correlation-weighted aggregation of support prototypes. It leverages superpixel-based pseudo-labels during training and introduces a quadrant masking scheme to suppress false positives in the downstream segmentation task. Key contributions include the correlation-weighted prototype aggregation module, per-pixel dynamic prototypes, and a domain-knowledge-informed quadrant masking strategy, validated on abdominal MR and CT datasets. Empirical results on CHAOS and SABS show competitive Dice scores against state-of-the-art methods without fine-tuning, highlighting robustness and practical impact for low-data medical imaging scenarios.

Abstract

Medical image segmentation is one of the domains where sufficient annotated data is not available. This necessitates the application of low-data frameworks like few-shot learning. Contemporary prototype-based frameworks often do not account for the variation in features within the support and query images, giving rise to a large variance in prototype alignment. In this work, we adopt a prototype-based self-supervised one-way one-shot learning framework using pseudo-labels generated from superpixels to learn the semantic segmentation task itself. We use a correlation-based probability score to generate a dynamic prototype for each query pixel from the bag of prototypes obtained from the support feature map. This weighting scheme helps to give a higher weightage to contextually related prototypes. We also propose a quadrant masking strategy in the downstream segmentation task by utilizing prior domain information to discard unwanted false positives. We present extensive experimentations and evaluations on abdominal CT and MR datasets to show that the proposed simple but potent framework performs at par with the state-of-the-art methods.

Correlation Weighted Prototype-based Self-Supervised One-Shot Segmentation of Medical Images

TL;DR

The paper tackles the challenge of medical image segmentation under limited annotations by proposing a self-supervised, one-shot framework that generates a dynamic prototype for every query pixel through correlation-weighted aggregation of support prototypes. It leverages superpixel-based pseudo-labels during training and introduces a quadrant masking scheme to suppress false positives in the downstream segmentation task. Key contributions include the correlation-weighted prototype aggregation module, per-pixel dynamic prototypes, and a domain-knowledge-informed quadrant masking strategy, validated on abdominal MR and CT datasets. Empirical results on CHAOS and SABS show competitive Dice scores against state-of-the-art methods without fine-tuning, highlighting robustness and practical impact for low-data medical imaging scenarios.

Abstract

Medical image segmentation is one of the domains where sufficient annotated data is not available. This necessitates the application of low-data frameworks like few-shot learning. Contemporary prototype-based frameworks often do not account for the variation in features within the support and query images, giving rise to a large variance in prototype alignment. In this work, we adopt a prototype-based self-supervised one-way one-shot learning framework using pseudo-labels generated from superpixels to learn the semantic segmentation task itself. We use a correlation-based probability score to generate a dynamic prototype for each query pixel from the bag of prototypes obtained from the support feature map. This weighting scheme helps to give a higher weightage to contextually related prototypes. We also propose a quadrant masking strategy in the downstream segmentation task by utilizing prior domain information to discard unwanted false positives. We present extensive experimentations and evaluations on abdominal CT and MR datasets to show that the proposed simple but potent framework performs at par with the state-of-the-art methods.
Paper Structure (35 sections, 14 equations, 3 figures, 5 tables)

This paper contains 35 sections, 14 equations, 3 figures, 5 tables.

Figures (3)

  • Figure 1: The figure depicts the entire working principle of the proposed framework. For clarity, we have also indicated the novel proposed correlation-weighted prototype aggregation step using a dotted red bounding box. T indicates the transformation applied to the support image to generate the query image only in the pre-training stage. Pool denotes pooling the feature map of the region denoted by the mask. MatMul denotes Matrix Multiplication. 'EM+SOC' denotes Element-wise Multiplication and Sum over Channels. Concat denotes the concatenation operation. (Best viewed at 300%)
  • Figure 2: Predictions in training phase at 25K, 50K, 75K, 100K iterations. The left image in Figs. \ref{['fig:tr1']}-\ref{['fig:tr4']} is the support image $\mathcal{X}_s$ and the support mask is denoted in green. The right image in Figs. \ref{['fig:tr1']}-\ref{['fig:tr4']} is the query image. The ground truth is denoted by green and the predicted mask is indicated by red. (Use 300% zoom for better visibility)
  • Figure 3: Figure showing the predictions obtained for 4 organs, Right Kidney, Left Kidney, Liver, and Spleen for two different modalities MR (CHAOS dataset) and CT (SABS dataset). (green) Ground Truth, (red) Prediction, (yellow) Ground Truth and Prediction overlap. (Use 300% zoom for better visibility)