Source-linked AI summary
Multi-organ Segmentation over Partially Labeled Datasets with Multi-scale Feature Abstraction
Xi Fang, Pingkun Yan
TL;DR
Limited fully annotated data makes multi-organ segmentation difficult, especially when useful datasets contain only partial labels. The paper introduces PIPO-FAN with unified training across partially labeled datasets and hierarchical multi-scale feature abstraction, reporting competitive performance on four public datasets. Its conclusion is that combining these mechanisms supports robust multi-organ segmentation and broader use of heterogeneous CT datasets.
Problem
Fully annotated datasets are scarce, and multi-organ segmentation is especially difficult because available datasets differ in their annotations.
Method
PIPO-FAN combines pyramid input and output processing, equal convolutional depth, deep supervision, adaptive multi-scale fusion, and unified training over partially labeled datasets.
Results
The method achieved very promising or competitive performance on BTCV, LiTS, KiTS, and Spleen, outperforming other networks on BTCV and the combined datasets.
Takeaways & Limitations
The unified strategy can use heterogeneous partially annotated datasets to train one network for multi-organ CT segmentation.
Abstract
from arXiv · showhide
Shortage of fully annotated datasets has been a limiting factor in developing deep learning based image segmentation algorithms and the problem becomes more pronounced in multi-organ segmentation. In this paper, we propose a unified training strategy that enables a novel multi-scale deep neural network to be trained on multiple partially labeled datasets for multi-organ segmentation. In addition, a new network architecture for multi-scale feature abstraction is proposed to integrate pyramid input and feature analysis into a U-shape pyramid structure. To bridge the semantic gap caused by directly merging features from different scales, an equal convolutional depth mechanism is introduced. Furthermore, we employ a deep supervision mechanism to refine the outputs in different scales. To fully leverage the segmentation features from all the scales, we design an adaptive weighting layer to fuse the outputs in an automatic fashion. All these mechanisms together are integrated into a Pyramid Input Pyramid Output Feature Abstraction Network (PIPO-FAN). Our proposed method was evaluated on four publicly available datasets, including BTCV, LiTS, KiTS and Spleen, where very promising performance has been achieved. The source code of this work is publicly shared at https://github.com/DIAL-RPI/PIPO-FAN for others to easily reproduce the work and build their own models with the introduced mechanisms.
I. INTRODUCTION
The paper addresses multi-organ segmentation with limited and partially labeled data by combining a unified training strategy with a new multi-scale PIPO-FAN architecture. The method uses hierarchical feature integration, equal convolutional depth, adaptive fusion, and deep supervision, and reports competitive results across multiple datasets.
- Related approaches: Existing segmentation methods commonly use encoder-decoder or pyramid structures, including U-Net and pyramid parsing modules, to exploit multi-scale features.P-IA parses input pyramids, whereas P-FA parses intermediate features and often combines scales late in the network.
- Proposed approach: PIPO-FAN integrates pyramid inputs and outputs with dedicated convolutional paths to maintain hierarchical information and reduce semantic gaps between scales.The architecture uses equal convolutional depth for fused features and an attention-based adaptive weight layer for multi-scale fusion.
- Motivation: Multi-organ segmentation is important for clinical applications, but collecting comprehensive multi-organ annotations is difficult.The paper motivates using related datasets created for different purposes despite their annotation differences.
- Evaluation: The method was evaluated on BTCV, LiTS, KiTS, and spleen datasets, achieving very promising performance compared with state-of-the-art approaches.The paper presents the evaluation as evidence of competitive multi-organ segmentation performance across publicly available datasets.
- Proposed approach: A unified training strategy enables one model to learn multi-organ segmentation from multiple partially labeled datasets.The contributions include a target-adaptive loss for handling partial annotations.
II. RELATED WORKS
Related work uses skip connections or pyramid parsing to exploit multi-scale context in segmentation. The paper identifies semantic gaps when features from different abstraction levels are fused and motivates combining these strategies more carefully for multi-organ segmentation.
- Skip-net methods: Skip-net methods combine features from different network levels through skip connections in pyramid structures such as U-Net and FPN.These features provide multi-scale context because receptive fields increase through the network.
- Pyramid parsing: Pyramid parsing methods extract multi-scale features at a common convolutional level using pyramid input analysis or pyramid feature analysis.The resulting features have different effective receptive fields and are concatenated or summed to strengthen contextual representation.
- Research gap: Pyramid inputs and parsing modules had not been explored together for multi-organ segmentation, although both can extract multi-scale context.The paper proposes combining them to use more of the pyramid structure in U-Net-like architectures.
B. Segmentation over multiple datasets
Training one segmentation model across datasets is difficult when annotations differ, even when images share related appearance information. The paper motivates partial-label learning as a way to combine complementary datasets for broader multi-organ training.
- Multiple-dataset learning: A general model trained over multiple datasets could learn more robust and accurate features, but differing annotations make direct generalization challenging.Direct multi-dataset training is straightforward only when label sets are the same.
- Multi-organ segmentation: Existing multi-organ segmentation methods are often organ-specific and require professional anatomical knowledge or manual design.Deep CNN methods broaden the approach but have frequently focused on single-organ segmentation.
- Architectural context: The proposed PIPO architecture propagates information from pyramid inputs through hierarchical abstraction and combination toward pyramid outputs.This design is presented as part of the paper’s response to multi-scale learning needs in segmentation.
- Partial-label learning: Partial annotations from complementary datasets provide a strategy for using available labels to train a unified segmentation model.The paper treats shared labeled organs across datasets as an opportunity to relax the requirement for fully annotated training data.
III. PYRAMID INPUT AND PYRAMID OUTPUT FEATURE ABSTRACTION NETWORK
PIPO-FAN integrates pyramid input and output feature abstraction in a U-shaped network, using equal convolutional depth and deep pyramid supervision to maintain multi-scale contextual information and refine outputs.
- Network architecture: PIPO-FAN integrates pyramid input analysis and pyramid feature analysis within a single network to fuse multi-scale context and semantically similar features.The network performs spatial pyramid pooling on the input and hierarchical multi-scale feature abstraction at each level.
- Pyramid input: Pyramid analysis uses spatial pyramid pooling and shared convolution to obtain multi-scale context at every network depth, combining global and local information.Unlike classical U-Net designs, PIPO-FAN maintains multi-scale features at each depth before fusing them into hierarchical representations.
- Equal convolutional depth: Equal convolutional depth applies the same number of convolutional layers to features fused at each level, reducing semantic gaps between scales.ResBlocks implement ECD so fused features share the same semantic abstraction level.
- Adaptive fusion: Adaptive fusion combines multi-scale segmentation outputs, whose lower scales provide local detail while higher scales provide class information despite blur.Adaptive weights are computed from the pyramid output features using a shared convolutional module.
- Pyramid output: Deep pyramid supervision generates labels at all output scales and computes scale-matched weighted cross-entropy losses during training.The mechanism supervises outputs with ground-truth segmentations spatially pooled to the corresponding scales.
- Pyramid output: Deep pyramid supervision can relieve gradient vanishing, support learning of deep features with hierarchical context, and preserve structural information across output scales.These effects follow from supervising outputs at multiple scales.
C. Adaptive Fusion
The adaptive fusion module learns the relative importance of segmentation outputs at different scales and automatically combines them into a final prediction.
- Fusion mechanism: Adaptive fusion learns the relative importance of each scale and automatically fuses multi-scale score maps before final segmentation.The module operates on score maps, defined as the last-layer outputs before softmax.
- Fusion mechanism: Pyramid output features are passed through an attention module after hierarchical abstraction from the pyramid input.The attention mechanism assigns scale-specific importance.
- Scale weighting: A shared convolutional block processes pyramid outputs to exploit similar structural information across scales and achieve scale invariance.Sharing the convolutional block provides a common transformation for outputs from different scales.
- Scale weighting: Global average pooling and global max pooling are summed to extract a global certainty score for each scale after squeezing features into single-channel vectors.The resulting score represents the overall output score at that scale.
- Output fusion: Softmax converts concatenated scale scores into weights, which rescale and sum pyramid outputs at the original image size before a final softmax segmentation.The weights reflect the importance of features at each scale.
D. Multi-organ segmentation over multiple datasets
The target adaptive loss enables unified multi-organ segmentation training across datasets with different partial labels by preserving labeled-organ probabilities and merging the others.
- Target adaptive loss enables training one segmentation algorithm across multiple datasets with different labels.It supports unified learning when datasets annotate different organ subsets.
- Partial-label probabilities are merged according to known labels, while unknown labels are treated as background for loss computation.
- The network uses multiple output branches, preserving labeled-organ probabilities and merging other branches into a dataset-specific non-target class.For LiTS, liver probability is preserved while all other organ probabilities form non-liver.
- Gradients backpropagate through branches associated with labeled organs, producing a general model that segments multiple organs when present.
A. Materials
Experiments use four public CT datasets with differing annotation scopes, standardized preprocessing, multi-scale training settings, and Dice-based evaluation with five-fold validation for LiTS and KiTS.
- Datasets: The evaluation combines LiTS, KiTS, Spleen, and BTCV, with single-organ annotations in the first three and multi-organ annotations in BTCV.
- Datasets: LiTS contains 131 training and 70 test images, while KiTS contains 210 training and 90 test images from 300 patients.
- Preprocessing: Axial slices are resized from 512×512 to 256×256 pixels, HU values are clipped to [-200, 200], and volumes are normalized to zero mean and unit variance.
- Implementation: The network uses five input and output scales, with convolutional feature-map channels progressing through the specified encoder-decoder blocks.
- Training: Deep supervision is used during the first 2000 epochs before adaptive weighting is activated for optimizing the fused score map.
- Evaluation: Performance is evaluated with Dice score, using five-fold cross validation on LiTS and KiTS and majority voting for LiTS challenge submissions.
C. Multi-organ Segmentation
The unified training strategy improves multi-organ segmentation by combining partially labeled datasets, while PIPO-FAN outperforms benchmark networks across BTCV and combined datasets. Ablations support multi-scale inputs and outputs, equal convolutional depth, and deep supervision as important design choices.
- Dataset combinations: Adding LiTS, KiTS, or Spleen data to BTCV improves multi-organ segmentation, with the largest spleen gains from additional datasets.The authors attribute the spleen improvement to similar appearances across datasets.
- Benchmark comparison: PIPO-FAN outperforms U-Net, ResU-Net, DenseU-Net, and DeepLabV3 on BTCV and the combined multi-dataset evaluation.The comparisons use benchmark networks trained under the same experimental setting.
- Benchmark comparison: Five-fold cross-validation shows PIPO-FAN significantly outperforms U-Net, ResU-Net, and DenseU-Net on LiTS, with p-values of 0.004, 0.025, and 0.002, respectively.The 131 labeled LiTS and 210 KiTS cases were split into five folds, with each fold used once for validation.
- Ablation studies: PIPO multi-scale input and output consistently outperforms single-scale segmentation, while equal convolutional depth further improves performance.ECD with deep pyramid supervision performs better than adaptive fusion without ECD.
- Ablation studies: Increasing both input and output scale counts improves accuracy, but increasing input scales alone from three to five reduces performance.The final model uses five input and five output scales, based on these experiments.
E. Comparison with state-of-the-arts on LiTS challenge
Two-step liver segmentation methods can be computationally expensive. The proposed single-step 2D method achieves competitive LiTS performance while reducing training and inference time.
- Computational cost: Two-step liver segmentation methods use coarse localization followed by fine segmentation, but can require substantial computation.One cited method required 21 hours and another 9 hours of fine-tuning using two Titan Xp GPUs.
- Computational cost: The proposed method trains on one Titan Xp GPU in 3 hours and segments one CT slice in 0.04s on a single GPU.The authors report this as faster than other reported methods.
- Segmentation performance: PIPO-FAN obtains the same Dice performance and competitive surface-distance scores, with less than a 0.2% Dice drop versus DeepX.Reported surface distances include ASSD 1.413 versus 1.450, MSSD 24.408 versus 27.118, and 2.421 versus 3.150.
V. DISCUSSION AND CONCLUSION
The paper presents a multi-scale network and unified training strategy for multi-organ segmentation with partially annotated datasets. The authors report competitive single-step performance and improved robustness and accuracy from enlarging the training dataset.
- Implications: The approach is complementary to semi-supervised learning and can be combined with it to further address data hunger in multi-organ segmentation.The paper also open-sources the implementation for further testing and development.
- Contributions: PIPO-FAN hierarchically incorporates multi-scale features at different network depths for medical image segmentation.The architecture is designed for multi-scale feature abstraction.
- Contributions: The unified strategy trains one segmentation network on multiple partially labeled datasets, improving robustness and accuracy by enlarging the training dataset.The strategy supports multi-organ segmentation across datasets with incomplete annotations.
- Implications: The 2D single-step network shows competitive CT segmentation performance compared with multi-step 3D networks.The authors state that the method can be extended to heterogeneous datasets covering abdominal, chest, pelvic, and whole-body regions.