[Paper Review] A two-stage 3D Unet framework for multi-class segmentation on full resolution image

Chengjia Wang, Tom MacGillivray+3 more|arXiv (Cornell University)|Apr 12, 2018

Advanced Image Processing Techniques16 references60 citations

TL;DR

This paper presents a two-stage concatenated U-Net framework that performs multi-class segmentation on full-resolution 3D medical images by first localizing a ROI and then refining segmentation without resampling, achieving superior performance to state-of-the-art U-Nets on MM-WHS data.

ABSTRACT

Deep convolutional neural networks (CNNs) have been intensively used for multi-class segmentation of data from different modalities and achieved state-of-the-art performances. However, a common problem when dealing with large, high resolution 3D data is that the volumes input into the deep CNNs has to be either cropped or downsampled due to limited memory capacity of computing devices. These operations lead to loss of resolution and increment of class imbalance in the input data batches, which can downgrade the performances of segmentation algorithms. Inspired by the architecture of image super-resolution CNN (SRCNN) and self-normalization network (SNN), we developed a two-stage modified Unet framework that simultaneously learns to detect a ROI within the full volume and to classify voxels without losing the original resolution. Experiments on a variety of multi-modal volumes demonstrated that, when trained with a simply weighted dice coefficients and our customized learning procedure, this framework shows better segmentation performances than state-of-the-art Deep CNNs with advanced similarity metrics.

Motivation & Objective

• Motivate high-resolution 3D whole-heart segmentation without losing detail through downsampling or post-processing.
• Develop a two-stage DCNN that first localizes a ROI and then refines voxel-wise labels at full resolution.
• Leverage a proximal SRCNN-inspired refinement with skip connections to improve segmentation accuracy.
• Demonstrate effectiveness on multi-modal cardiac CT/MR data from the MM-WHS challenge.
• Compare with state-of-the-art 3D U-Nets and assess training strategies to achieve good performance with limited data.

Proposed method

• A two concatenated U-Net architecture where Net1 provides coarse, ROI-focused predictions on down-sampled volumes and Net2 refines labels at full resolution.
• Net1 uses dilated 5x5x5 convolutions with increasing dilation in the contracting path to enlarge the effective receptive field.
• Net2 follows an SRCNN-inspired design with skip connections and recursive-like structure, taking as input both Net1 output and the original data.
• A 4-step training procedure includes: (1) ROI localization with full volumes, (2) coarse multi-class segmentation with partial volumes, (3) joint coarse+fine segmentation with end-to-end training, (4) fine-tuning Net2 on subvolumes of axial slices (K slices).
• A weighted multi-class Dice-based loss (Eq. 2) plus a foreground-focused Dice loss for ROI localization (Eq. 3–4) guide training; an Adam optimizer is used with a learning rate of 0.0001.
• Data augmentation and hierarchical/subvolume sampling address data imbalance and memory constraints; batch design focuses on central voxel class to balance low-frequency classes.

Experimental results

Research questions

• RQ1 Can a two-stage, full-resolution 3D segmentation framework outperform single-shot 3D U-Nets without post-processing or resampling?
• RQ2 Does ROI-based localization combined with a refinement stage improve multi-class segmentation in high-resolution cardiac CT/MR data?
• RQ3 How do step-wise training procedures and specialized loss functions affect convergence and accuracy on MM-WHS data?
• RQ4 Is the proposed framework competitive with state-of-the-art U-Nets on MM-WHS benchmarks for both CT and MR modalities?

Key findings

• Net2 achieved higher Dice and Jaccard scores than a baseline 3D U-Net across CT structures.
• Net1 provided improvements over a downsampled U-Net on several structures, especially when used in conjunction with Net2.
• On MM-WHS test data, the framework yielded very competitive MR results (average Dice around 0.83 for MR), and CT results showed notable gains over baselines; the results approached the competition winner with fewer training iterations and without additional post-processing.
• The method demonstrates effective segmentation at full-resolution, validating the ROI-driven, two-stage approach for high-resolution 3D medical image segmentation.
• End-to-end training combining coarse and fine losses accelerates convergence and improves performance compared to training separate networks.