[Paper Review] CFUN: Combining Faster R-CNN and U-net Network for Efficient Whole Heart Segmentation

Zhanwei Xu, Ziyi Wu+1 more|arXiv (Cornell University)|Dec 12, 2018

Advanced Neural Network Applications16 references46 citations

TL;DR

CFUN integrates a modified 3D Faster R-CNN for heart localization with a 3D U-net for segmentation, augmented by a 3D Edge-loss head, yielding fast and competitive whole-heart segmentation.

ABSTRACT

In this paper, we propose a novel heart segmentation pipeline Combining Faster R-CNN and U-net Network (CFUN). Due to Faster R-CNN's precise localization ability and U-net's powerful segmentation ability, CFUN needs only one-step detection and segmentation inference to get the whole heart segmentation result, obtaining good results with significantly reduced computational cost. Besides, CFUN adopts a new loss function based on edge information named 3D Edge-loss as an auxiliary loss to accelerate the convergence of training and improve the segmentation results. Extensive experiments on the public dataset show that CFUN exhibits competitive segmentation performance in a sharply reduced inference time. Our source code and the model are publicly available at https://github.com/Wuziyi616/CFUN.

Motivation & Objective

• Motivate efficient whole heart segmentation without patch-based tiling by focusing on the heart region.
• Develop a detection-segmentation pipeline by embedding 3D U-net into a modified Faster R-CNN to segment all heart structures in one step.
• Introduce a 3D Edge-loss head to accelerate training and improve boundary accuracy of segmentation.
• Evaluate CFUN on MM-WHS2017 and compare with state-of-the-art methods in terms of Dice score and inference time.

Proposed method

• Extend Faster R-CNN to 3D with a P3D Bottleneck to handle CT volume dimensions.
• Replace ResNet with P3D and add FPN for multi-scale feature fusion in the RPN for heart localization.
• Embed a modified 3D U-net as the segmentation network that takes the RoI-aligned heart region as input.
• Use deep supervision in the decoder and an extra deconvolution layer to double output size for finer segmentation.
• Introduce a 3D Edge-loss head using fixed 3D Sobel-kernels to enforce edge consistency between prediction and ground truth.
• Train with a composite loss: L = w1 L_box + w2 L_cls + w3 L_seg + w4 L_edge with weights 2:2:2:1.

Experimental results

Research questions

• RQ1 Can a detection-based pipeline outperform patch-based tiling approaches in whole heart segmentation?
• RQ2 Does embedding a 3D U-net inside a modified Faster R-CNN improve segmentation accuracy and reduce inference time?
• RQ3 Does the 3D Edge-loss head contribute to boundary accuracy and faster convergence?
• RQ4 How does CFUN compare to existing 3D U-net and multi-stage CNN approaches on MM-WHS2017 in Dice score and speed?

Key findings

• CFUN achieves an average Dice score of 0.859 on MM-WHS2017 test data, competitive with the competition winner.
• CFUN segments a heart in under 10 seconds (under 15 seconds including data loading and post-processing).
• Ablation shows the bounding box refine head improves detection precision and the Edge-loss head enhances boundary quality.
• Compared to 3D U-net and two-stage U-nets, CFUN delivers faster inference with comparable or better performance.
• CFUN’s single bounding box approach avoids tiling and patch-wise processing, leveraging global heart information for regularized segmentation.