[Paper Review] A New Probabilistic V-Net Model with Hierarchical Spatial Feature Transform for Efficient Abdominal Multi-Organ Segmentation

Minfeng Xu, Heng Guo+3 more|arXiv (Cornell University)|Aug 2, 2022

Radiomics and Machine Learning in Medical Imaging4 citations

TL;DR

This paper proposes a probabilistic V-Net with hierarchical spatial feature transformation for efficient and accurate abdominal multi-organ segmentation. By leveraging a conditional variational autoencoder to model organ-specific semantic variations in latent space and injecting them via spatial-wise affine transformations, the method improves Dice scores by 7.3% for kidneys and 9.7% for pancreas while achieving ~7× inference speedup over nnUNet and CoTr.

ABSTRACT

Accurate and robust abdominal multi-organ segmentation from CT imaging of different modalities is a challenging task due to complex inter- and intra-organ shape and appearance variations among abdominal organs. In this paper, we propose a probabilistic multi-organ segmentation network with hierarchical spatial-wise feature modulation to capture flexible organ semantic variants and inject the learnt variants into different scales of feature maps for guiding segmentation. More specifically, we design an input decomposition module via a conditional variational auto-encoder to learn organ-specific distributions on the low dimensional latent space and model richer organ semantic variations that is conditioned on input images.Then by integrating these learned variations into the V-Net decoder hierarchically via spatial feature transformation, which has the ability to convert the variations into conditional Affine transformation parameters for spatial-wise feature maps modulating and guiding the fine-scale segmentation. The proposed method is trained on the publicly available AbdomenCT-1K dataset and evaluated on two other open datasets, i.e., 100 challenging/pathological testing patient cases from AbdomenCT-1K fully-supervised abdominal organ segmentation benchmark and 90 cases from TCIA+&BTCV dataset. Highly competitive or superior quantitative segmentation results have been achieved using these datasets for four abdominal organs of liver, kidney, spleen and pancreas with reported Dice scores improved by 7.3% for kidneys and 9.7% for pancreas, while being ~7 times faster than two strong baseline segmentation methods(nnUNet and CoTr).

Motivation & Objective

• To address the challenge of high variability in abdominal organ shapes and appearances across CT scans.
• To improve segmentation robustness and accuracy for liver, kidney, spleen, and pancreas in diverse anatomical and pathological cases.
• To develop a fast inference method that maintains high performance on multi-organ segmentation tasks.
• To model complex organ-specific semantic variations using a learnable latent space conditioned on input images.
• To integrate learned variations hierarchically into feature maps for improved spatial guidance in segmentation.

Proposed method

• A conditional variational autoencoder (CVAE) is used to learn organ-specific distributions in a low-dimensional latent space from input CT images.
• The latent codes from the CVAE are used to predict spatial-wise affine transformation parameters for modulating feature maps at multiple scales.
• Hierarchical spatial feature transformation injects learned semantic variations into decoder features at different levels, enhancing fine-grained segmentation.
• The V-Net architecture is extended with probabilistic decoding to model uncertainty and improve robustness to shape and intensity variations.
• The model is trained end-to-end on the AbdomenCT-1K dataset and evaluated on two open benchmarks: AbdomenCT-1K (100 test cases) and TCIA+&BTCV (90 cases).
• The method uses spatial-wise affine transformations to conditionally modulate feature maps, enabling dynamic adaptation to organ-specific patterns.

Experimental results

Research questions

• RQ1 Can a probabilistic V-Net with learned latent organ variations improve segmentation accuracy on challenging abdominal CT scans?
• RQ2 How does hierarchical spatial feature transformation enhance multi-scale feature modulation for organ segmentation?
• RQ3 Can the integration of organ-specific semantic variations lead to faster inference without sacrificing performance?
• RQ4 To what extent does the model generalize to pathological and diverse anatomical cases?
• RQ5 How does the method compare in performance and speed to state-of-the-art models like nnUNet and CoTr?

Key findings

• The proposed method achieved a 7.3% relative improvement in Dice score for kidneys compared to strong baselines on the AbdomenCT-1K test set.
• A 9.7% Dice score improvement was recorded for the pancreas, indicating strong gains on the most challenging organ.
• The model achieved ~7 times faster inference speed than both nnUNet and CoTr, significantly improving efficiency.
• The method demonstrated robust performance on 100 challenging and pathological cases from AbdomenCT-1K, confirming generalization capability.
• The integration of CVAE-learned organ-specific variations via spatial-wise affine transformations led to improved feature modulation and segmentation accuracy.
• Quantitative results on the TCIA+&BTCV dataset further validated the method’s superior performance and generalization across diverse data distributions.