Source-linked AI summary

ShapePU: A New PU Learning Framework Regularized by Global Consistency for Scribble Supervised Cardiac Segmentation

Ke Zhang, Xiahai Zhuang

arXiv:2206.02118v1cs.CV

TL;DR

Dense pixel-wise labels are costly, and scribble supervision leaves insufficient labels and incomplete shape information. ShapePU combines PU learning with global consistency regularization, achieving state-of-the-art performance across two cardiac segmentation datasets.

  • Problem

    Pixel-wise dense labeling is costly, while scribble-guided segmentation faces insufficient supervision and incomplete object shapes.

    Source: §1 · pp. 1–2
  • Method

    ShapePU uses EM to estimate unlabeled-pixel class ratios, marginal probability maximization to assign their classes, and cutout-based global consistency to regularize shape.

    Source: §1 · p. 2, §2.2 · p. 4, §4 · p. 9
  • Results

    ShapePU consistently outperformed fully supervised, state-of-the-art weakly supervised, and competitive PU learning methods on two cardiac segmentation tasks.

    Source: §1 · p. 3, §4 · p. 9
  • Takeaways & Limitations

    ShapePU trains using scribble-annotated data alone, without requiring fully annotated masks for training.

    Source: §1 · p. 2

    Mixture-proportion estimation assumes labeled and unlabeled pixels share the same within-class probability distribution.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

Cardiac segmentation is an essential step for the diagnosis of cardiovascular diseases. However, pixel-wise dense labeling is both costly and time-consuming. Scribble, as a form of sparse annotation, is more accessible than full annotations. However, it's particularly challenging to train a segmentation network with weak supervision from scribbles. To tackle this problem, we propose a new scribble-guided method for cardiac segmentation, based on the Positive-Unlabeled (PU) learning framework and global consistency regularization, and termed as ShapePU. To leverage unlabeled pixels via PU learning, we first present an Expectation-Maximization (EM) algorithm to estimate the proportion of each class in the unlabeled pixels. Given the estimated ratios, we then introduce the marginal probability maximization to identify the classes of unlabeled pixels. To exploit shape knowledge, we apply cutout operations to training images, and penalize the inconsistent segmentation results. Evaluated on two open datasets, i.e, ACDC and MSCMRseg, our scribble-supervised ShapePU surpassed the fully supervised approach respectively by 1.4% and 9.8% in average Dice, and outperformed the state-of-the-art weakly supervised and PU learning methods by large margins. Our code is available at https://github.com/BWGZK/ShapePU.

1 Introduction

Dense pixel labels are costly, while scribble supervision is sparse and difficult to use because unlabeled pixels and complete object shape remain challenging. ShapePU combines multi-class PU learning with cutout-based global consistency, and the authors report it outperformed fully supervised, weakly supervised, and PU methods on two cardiac segmentation tasks.

1 Introduction

  • Dense pixel-wise labeling is burdensome, motivating interest in more accessible sparse annotations such as scribbles.Source: §1 · p. 1
  • Scribble-guided segmentation faces insufficient supervision and incomplete object shape, while prior shape-learning approaches may require additional fully annotated masks.Source: §1 · p. 2
  • Existing PU methods estimate each class's mixture ratio independently, which the authors consider unreasonable for multi-class segmentation.Source: §1 · p. 2
  • ShapePU estimates multi-class mixture ratios with EM, assigns unlabeled pixels by marginal-probability maximization, and regularizes shape through cutout equivalence.Source: §1 · p. 2
  • The authors describe ShapePU as the first PU framework for weakly supervised segmentation incorporating shape knowledge.Source: §1 · p. 2
  • ShapePU consistently outperformed fully supervised, state-of-the-art weakly supervised, and competitive PU methods on two cardiac segmentation tasks.Source: §1 · p. 3

2 Method

ShapePU combines scribble-supervised learning with a multi-class PU framework for using unlabeled pixels and global consistency regularization for shape. Its pipeline estimates class proportions, assigns unlabeled pixels through marginal probabilities, and regularizes predictions on cutout images.

2 Method

  • ShapePU combines supervised learning on annotated pixels with multi-class PU learning and cutout-based global consistency to use unlabeled pixels and regularize shape.The three components are EM-based mixture-proportion estimation, marginal-probability PU learning, and cutout equivalence.Source: §2 · p. 3

2.1 EM Estimation of Multi-class Mixture Proportion

  • EM estimates class proportions in the unlabeled-pixel mixture, initializing them from labeled-pixel frequencies and iteratively updating until the estimates converge.The procedure assumes labeled and unlabeled pixels share the same within-class probability.Source: §2.1 · pp. 3–4

2.2 PU Learning with Marginal Probability

  • For each foreground class, ShapePU ranks unlabeled pixels by predicted class probability and uses the estimated proportion to define positive and negative sets.Because foreground-class ratios tend to be overestimated, the loss is applied to predicted negatives and maximizes the marginal probability of the fused other classes.Source: §2.2 · p. 4

2.3 Global Consistency

  • ShapePU cuts out a square region and applies rotation or flipping, requiring predictions on the original and transformed images to agree outside the cutout.The global-consistency loss uses cosine similarity distance in both prediction directions.Source: §2.3 · pp. 4–5
  • The training objective also applies supervised cross-entropy to scribble-labeled pixels in both the original and transformed images.The optimization objective combines this supervised term with the method’s other losses.Source: §2.3 · p. 5

3 Experiment

Experiments evaluate ShapePU on ACDC and MSCMRseg through an ablation and comparisons with weakly supervised and PU methods. ShapePU outperformed the compared methods, while ablations show gains from unlabeled-pixel supervision and global consistency.

3.1 Experiment Setup

  • Experiments used ACDC and MSCMRseg, with 70/15/15 training, validation, and test subjects on ACDC and 25/5/15 images on MSCMRseg.The setup also describes the datasets, scribble annotations, image preprocessing, and training configuration.Source: §3.1 · pp. 5–6, §3 · p. 5

3.2 Ablation study

  • A 1.7% average Dice gain (85.1% vs 83.4%) followed adding global consistency to PU with cutout; negative loss alone gained 5.5% (83.3% vs 77.8%).Cutout alone gave marginal improvement and could alter target shape; global consistency imposes cutout equivalence.Source: §3.2 · pp. 6–7, Table 1 · p. 6

3.3 Comparisons with weakly supervised methods

  • ShapePU outperformed the compared scribble-guided, GAN-based, and fully supervised methods, matching UNetF with significantly better LV HD (p=0.041).The comparison uses 35 scribbles; Table 2 provides the quantitative results.Source: Table 2 · p. 7, §3.3 · p. 8

3.4 Comparisons with PU learning methods

  • 85.2% average Dice (85.2% vs 83.7%) and 15.0 mm vs 41.7 mm average HD: ShapePU improved on competing PU methods on MSCMRseg.Without ground-truth mixture ratios, the basic PU framework still outperformed peers with at least 33.9% average Dice.Source: Table 3 · p. 8, §3.4 · p. 9

4 Conclusion

The conclusion presents ShapePU as a shape-constrained PU method for weakly supervised cardiac segmentation. It uses EM mixture-ratio estimation and PU learning for unlabeled pixels, and shape-consistency loss on cutout equivalence; evaluation on two public datasets achieved state-of-the-art performance.

4 Conclusion

  • ShapePU combines EM estimation of unlabeled-pixel mixture ratios with PU learning to distinguish their classes.The method is presented for weakly supervised cardiac segmentation.Source: §4 · p. 9
  • Shape-consistency loss regularizes cutout equivalence to capture the heart’s global shape.Source: §4 · p. 9
  • Evaluation on two publicly available datasets achieved state-of-the-art performance.Source: §4 · p. 9
Loading 2206.02118v1…