Source-linked AI summary
Shape-aware Semi-supervised 3D Semantic Segmentation for Medical Images
Shuailin Li, Chuyu Zhang, Xuming He
TL;DR
Medical segmentation requires costly pixel-wise annotations, and prior semi-supervised methods may neglect geometry or require aligned shape priors. This paper jointly predicts segmentation and SDMs, using adversarial consistency between labeled and unlabeled SDMs; on the Atrial Segmentation Challenge dataset, it outperforms state-of-the-art methods with high-quality global shapes.
Problem
Medical image segmentation requires expensive annotations, while existing semi-supervised methods can neglect geometric constraints or require aligned shape priors.
Method
A multi-task network jointly predicts semantic segmentation and signed distance maps, with adversarial loss enforcing consistency between labeled and unlabeled SDM predictions.
Results
The method outperforms state-of-the-art methods and generates object segmentations with high-quality global shapes on the Atrial Segmentation Challenge dataset.
Takeaways & Limitations
Geometric shape regularity provides a constraint for both segment prediction and network learning in semi-supervised 3D medical segmentation.
Takeaways & Limitations
The formulation assumes N labeled and M unlabeled data with N ≪ M and directly treats each semantic class separately for multi-class SDMs.
Abstract
from arXiv · showhide
Semi-supervised learning has attracted much attention in medical image segmentation due to challenges in acquiring pixel-wise image annotations, which is a crucial step for building high-performance deep learning methods. Most existing semi-supervised segmentation approaches either tend to neglect geometric constraint in object segments, leading to incomplete object coverage, or impose strong shape prior that requires extra alignment. In this work, we propose a novel shapeaware semi-supervised segmentation strategy to leverage abundant unlabeled data and to enforce a geometric shape constraint on the segmentation output. To achieve this, we develop a multi-task deep network that jointly predicts semantic segmentation and signed distance map(SDM) of object surfaces. During training, we introduce an adversarial loss between the predicted SDMs of labeled and unlabeled data so that our network is able to capture shape-aware features more effectively. Experiments on the Atrial Segmentation Challenge dataset show that our method outperforms current state-of-the-art approaches with improved shape estimation, which validates its efficacy. Code is available at https://github.com/kleinzcy/SASSnet.
1 Introduction
Medical image segmentation needs extensive annotations, while existing semi-supervised methods either neglect object geometry or require strongly aligned shape priors. The proposed approach uses SDM-based geometric representations and adversarial learning to improve shape-aware segmentation from labeled and unlabeled data.
- Motivation: Medical segmentation annotation is expensive, motivating semi-supervised learning with small labeled and many unlabeled datasets.Semi-supervised frameworks typically use a small labeled dataset alongside many unlabeled images.
- Limitations of Prior Work: Existing approaches either use generic regularization that can neglect geometric constraints or impose strong anatomical priors requiring properly aligned images.Alignment is difficult when objects vary substantially in pose or shape.
- Proposed Strategy: The proposed strategy introduces a flexible geometric representation to enforce global shape constraints while accommodating objects with varying poses or shapes.The representation is intended to capture each object class’s global shape more effectively.
- Method: A shared-backbone multi-task network jointly predicts semantic segmentation and signed distance maps of object surfaces.An SDM assigns each pixel its signed distance to the nearest target-object boundary and encodes shape and surface information.
- Method: Adversarial learning matches predicted SDM distributions for labeled and unlabeled data, providing an implicit shape prior for semi-supervised volumetric segmentation.The method combines a multi-task supervised loss with an adversarial loss on SDM predictions.
- Results: On the Atrial Segmentation Challenge dataset, the method outperforms state-of-the-art methods and produces high-quality global object shapes.The evaluation uses extensive comparisons with prior methods.
2 Method
The method uses a multi-task network to predict segmentation maps and signed distance maps, then regularizes unlabeled-data learning through adversarial consistency of shape representations.
- Segmentation Network: The network jointly predicts a 3D object mask and its signed distance map using separate output branches.The shared V-Net backbone contains a segmentation branch and an SDM branch.
- Shape-aware Semi-supervised Learning: Training assumes N labeled examples and M unlabeled examples, with N ≪ M, and derives ground-truth SDMs from labeled segmentation annotations.The labeled set contains input volumes, segmentation annotations, and corresponding ground-truth SDMs.
- Supervised Loss: The supervised loss combines Dice loss for segmentation with mean squared error for SDM prediction, weighted by coefficient α.The two terms are denoted Lseg and Lsdm, with α balancing their contributions.
- Adversarial Loss: An adversarial loss uses a discriminator to distinguish predicted SDMs from labeled and unlabeled data, enforcing consistency across the two sets.The discriminator receives the input volume and predicted SDM, while minimizing the adversarial loss encourages shape-aware features that generalize to unlabeled data.
- Overall Training Pipeline: The overall objective combines supervised and adversarial losses, and training alternates optimization of the segmentation network and discriminator.The adversarial term is weighted by β, which is increased gradually with an annealing warm-up strategy.
3 Experiments and Results
Experiments on the Left Atrium dataset compare SASSNet with semi-supervised baselines and isolate the contributions of its SDM and adversarial-loss modules. SASSNet achieves strong segmentation results, especially with limited labeled data, while preserving object shape.
- Dataset and setup: The Left Atrium dataset contains 100 3D GE-MRI scans, split into 80 training and 20 validation scans at 0.625 × 0.625 × 0.625mm3 resolution.
- Dataset and setup: Evaluation uses Dice, Jaccard, 95HD, and ASD, with the segmentation map evaluated after NMS removes isolated extraneous regions.
- Quantitative comparison: With 16 labeled images, SASSNet outperforms all compared semi-supervised networks at 89.54% Dice and 81.24% Jaccard.
- Quantitative comparison: With only 8 labeled images, SASSNet surpasses UA-MT by +2.56% Dice without NMS and +3.07% with NMS.
- Ablative study: Adding the SDM head improves Dice by 1.1%, while adding unlabeled data and adversarial loss improves it by 5.7%.
- Ablative study: SASSNet outperforms both mean-teacher comparisons on Dice and Jaccard while using a simpler architecture.
4 Conclusion
The paper presents a shape-aware semi-supervised approach for 3D medical scans that combines segmentation, SDM prediction, and SDM consistency learning. On the Atrial Segmentation Challenge dataset, it outperforms state-of-the-art methods and produces high-quality global shapes.
- The approach exploits target-object geometric regularity for segment prediction and network learning.
- A multi-task network jointly predicts semantic segmentation and object-surface SDMs.
- A semi-supervised loss enforces consistency between predicted SDMs for labeled and unlabeled data.
- Validation on the Atrial Segmentation Challenge dataset shows outperformance of state-of-the-art methods with high-quality global shapes.