Source-linked AI summary
Recurrent Mask Refinement for Few-Shot Medical Image Segmentation
Hao Tang, Xingwei Liu, Shanlin Sun, Xiangyi Yan, Xiaohui Xie
TL;DR
Medical segmentation networks often require extensive manual labels and can struggle to generalize to unseen classes, while medical boundaries can be ambiguous because foreground and background appearances overlap. RP-Net combines a correlation-based context relation encoder with recurrent mask refinement using a prototypical network. Across three abdominal datasets, it reports DSC gains over prior state-of-the-art methods of 16.32%, 8.45%, and 6.24%, respectively.
Problem
Medical segmentation networks typically need many manual labels and may generalize poorly to unseen classes, while similar tissue appearances and inhomogeneous backgrounds complicate foreground boundaries.
Source: §1 · p. 1Method
RP-Net uses correlation in a context relation encoder to model foreground–background features, then recurrently recomputes context and refines masks with a prototypical network.
Source: §1 · p. 2Results
Takeaways & Limitations
The authors state that CRE and recurrent refinement are generic modules that can also be integrated into other network types to enhance context relationship features.
Source: §5 · p. 8
Abstract
from arXiv · showhide
Although having achieved great success in medical image segmentation, deep convolutional neural networks usually require a large dataset with manual annotations for training and are difficult to generalize to unseen classes. Few-shot learning has the potential to address these challenges by learning new classes from only a few labeled examples. In this work, we propose a new framework for few-shot medical image segmentation based on prototypical networks. Our innovation lies in the design of two key modules: 1) a context relation encoder (CRE) that uses correlation to capture local relation features between foreground and background regions; and 2) a recurrent mask refinement module that repeatedly uses the CRE and a prototypical network to recapture the change of context relationship and refine the segmentation mask iteratively. Experiments on two abdomen CT datasets and an abdomen MRI dataset show the proposed method obtains substantial improvement over the state-of-the-art methods by an average of 16.32%, 8.45% and 6.24% in terms of DSC, respectively. Code is publicly available.
1. Introduction
Medical segmentation systems often need extensive expert annotations and may generalize poorly to unseen classes. RP-Net addresses few-shot segmentation with foreground–background context modeling and recurrent mask refinement, and reports gains over prior state of the art on three abdominal datasets.
1. Introduction
- Few-shot learning segments novel classes from annotated support images, but its application to medical image segmentation remains at an early stage.The approach responds to the cost and time of obtaining extensive expert annotations and to image-acquisition differences that challenge generalization.Source: §1 · p. 1
- Medical few-shot segmentation can be ambiguous because foreground regions are small and homogeneous while backgrounds are large, inhomogeneous, and may contain visually similar tissues.These conditions make local appearance and context important for identifying foreground–background boundaries.Source: §1 · pp. 1–2
- RP-Net uses a context relation encoder (CRE) to model foreground–background feature relationships, augment pixel features, and guide a prototypical network to predict masks.Correlation captures differences between foreground and background regions, with emphasis on context around object boundaries.Source: §1 · p. 2
- The recurrent module reuses CRE and a prototypical network with shared weights, recomputing context from each prior prediction to refine the mask iteratively.Each updated prediction changes the local context relationship used in the next refinement step.Source: §1 · p. 2
- 16.32% on ABD-110 [49], 8.45% on MIC-CAI15 Multi-Atlas Abdomen Labeling challenge, and 6.24% on ISBI 2019 Combined Healthy Abdominal Organ Segmentation Challenge are the reported DSC gains over the state-of-the-art framework.The evaluation covers two abdomen CT datasets and one abdomen MRI dataset.Source: §1 · p. 2
2. Related work
The related work contrasts annotation-intensive medical segmentation models with few-shot approaches that use support examples to segment new classes. RP-Net builds on prototype-based few-shot segmentation with context-relation features and recurrent mask refinement.
2.1. Medical image segmentation
- Conventional medical segmentation architectures, including U-Net variants, use abundant manual annotations and are limited when segmenting novel classes.U-Net fuses encoder and decoder features through lateral connections; cited variants include V-Net, Attention U-Net, and nnUNet.Source: §2.1 · p. 2
2.2. Few-shot learning
- For few-shot medical segmentation without retraining, RP-Net combines CRE and recurrent mask refinement to capture local feature and shape differences around foreground boundaries.This builds on prototype-based segmentation, which pools class-specific features from support masks and segments query images using distances to class prototypes.Source: §2.2 · pp. 2–3, Figure 1 · p. 3
3. Method
Few-shot medical image segmentation uses labeled support examples to segment query images from unseen classes. RP-Net extracts features, models foreground–background context, and iteratively refines masks with a prototypical network.
3.1. Problem definition
- In N-way K-shot segmentation, training and test classes are disjoint, and labeled support examples guide segmentation of query images.The paper notes that medical-image segmentation commonly uses 1-way 1-shot learning.Source: §3.1 · p. 3
3.2. Proposed method
- RP-Net combines feature extraction, context relation encoding, and iterative CRE–prototypical mask refinement in a differentiable, end-to-end trainable pipeline.Source: §3.2 · p. 3
3.2.1 Feature extraction
- An adapted U-Net extracts support and query features at one-quarter of image resolution after global affine alignment.The final two U-Net upsampling blocks are removed to save memory and computation.Source: §3.2.1 · p. 4
3.2.2 Context relation encoder (CRE)
- CRE correlates masked foreground and background features within a local neighborhood, then fuses the resulting context features with foreground features.The displacement limit is d=5; separating foreground and background makes the correlation sparse around object boundaries.Source: §3.2.2 · p. 4
3.2.3 Prototypical networks
- Prototypical segmentation averages support features within each class mask to form prototypes, then classifies query features using distances to those prototypes.Source: §3.2.3 · p. 4
3.2.4 Recurrent mask refinement
- At each recurrent step, CRE recomputes context features from the previous mask before the prototypical network predicts an updated mask.Weights are shared across iterations; inference uses 10 iterations, which the paper reports are sufficient for a stable result.Source: §3.2.4 · pp. 4–5
3.2.5 Loss function
- The final predicted mask is supervised with a weighted sum of Dice loss and cross-entropy, with β set to 1.Source: §3.2.5 · p. 5
4. Experiment
Experiments compare RP-Net with prior methods across three abdominal datasets and examine its components, initialization, backbones, and iterative refinement. RP-Net outperformed SSL-ALPNet on all three datasets, while ablations assess the contributions of CRE and recurrence.
4.1. Setup
- Experiments used ABD-30 and ABD-MR, with five-fold cross-validation under 1-way 1-shot learning; liver, spleen, and left and right kidney were semantic classes, and overlap was evaluated using DSC.Source: §4.1 · p. 5
4.2. Comparison with the state-of-the-art methods
- RP-Net outperformed SSL-ALPNet by an average of 16.32%, 9.46%, and 6.24% on ABD-110, ABD-30, and ABD-MR, respectively.Source: Table 1 · p. 6, §4.2 · p. 7
4.3. Ablation study
- Affine + Concat outperforms Affine by 3.04%, while Affine + CRE outperforms Affine + Concat by 3.82%, showing CRE better captures local differences via correlation.Source: §4.3 · p. 7
- Using recurrent mask refinement, Affine + Concat + Recurrent performs 7.84% better than Affine + Concat; combining Affine + CRE + Recurrent yields a 15.39% improvement over the recurrent Affine + Concat variant.Source: §4.3 · p. 7
- U-Net exceeded Res18 by an average of 2.32%, while performance was not very sensitive to correlation radius d (0, 1, 3, 5, 7).The authors attribute U-Net’s gain mainly to lateral connections fusing low- and high-level features.Source: §4.3 · p. 7
- With recurrent training, DSC quickly converged across iterations; without it, performance diverged after iteration one.Source: Figure 2 · p. 8, §4.3 · p. 8
- Demons initialization improved the result by 1.02%, while RP-Net outperformed Grabcut by 19%; affine initialization sufficed when foreground was roughly located.Source: §4.3 · pp. 7–8
4.4. Qualitative result
- Across CT and MR examples, RP-Net gradually refined the initial mask and converged to a solution with clear object boundaries.Source: Figure 3 · p. 8, §4.4 · p. 8
5. Conclusion
RP-Net iteratively refines segmentation masks using a context relation encoder and recurrent module. Across three organ-segmentation datasets, it outperforms the previous state-of-the-art approach by as much as 16% in DSC.
5. Conclusion
- Up to 16% in DSC: RP-Net outperforms the previous state-of-the-art approach across three organ-segmentation datasets.The conclusion reports the maximum improvement across the evaluated datasets.Source: §5 · p. 8
- RP-Net repeatedly applies a context relation encoder and recurrent module to incrementally refine masks toward better object-boundary alignment.The conclusion identifies iterative mask refinement and boundary alignment as the framework’s central operation.Source: §5 · p. 8
- The context relation encoder and recurrent module are generic and can be integrated into other network types to enhance context-relation features.Source: §5 · p. 8