论文分享 CFUN: Combining Faster R-CNN and U-net Network for Efficient Whole Heart Segmentation

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

摘要

本文提出了一种结合Faster R-CNN和U-net网络的心脏分割方法。由于R-CNN具有更快的精确定位能力和U-net强大的分割能力，CFUN只需要一步检测和分割推理就可以得到完整的心脏分割结果，在显著降低计算成本的情况下获得了良好的分割效果。此外，CFUN采用了一种新的基于边缘信息的损失函数3D-edge-loss作为辅助损失，加快了训练的收敛速度，提高了分割效果。在公共数据集上的大量实验表明，CFUN在显著缩短的推理时间内表现出竞争性的分割性能。我们的源代码和模型在https://github.com/Wuziyi616/CFUN上公开发布。

介绍

由于GPU的内存限制，通常采用一种分块策略，即分割网络需要对体积较小的块进行处理，并将它们叠加在一起，形成最终的心脏分割结果。由于心脏体积只占整个CT图像的一小部分，这种方法无法捕捉心脏解剖的全局约束，浪费了大量的计算资源。

我们提出了一种基于检测的分割方法，使得分割网络能够聚焦于CT图像中最相关的部分，而不需要将整个图像分割成小块。另一个优点是分割网络以RoI所包围的整个心脏图像作为输入，可以利用心脏的整体信息得到更加规则化的分割结果。

与传统Faster R-CNN处理多个不同目标的boxes不同，我们迫使检测网络只识别一个包括了整个心脏标注bounding box。考虑到CT图像的成像特点，将Faster R-CNN中的ResNet结构替换为P3D Bottleneck，并选择3D Edge-loss分支作为3D U-net的辅助任务，以获得更精确的分割。

Faster R-CNN and Region-proposal Based Segmentation

经过两次改进，Faster R-CNN显示出卓越的检测和定位能力，并且非常有效。图片首先通过区域建议网络（RPN）得到多个推荐框，然后直接在调整后的推荐框对应特征中加入分类和回归网络，得到对象的分类结果和最终的边界框。通过在RPN后添加mask head，Mask R-CNN网络可以同时进行检测、分类和分割。目前已有一些尝试将Mask-R-CNN应用于医学图像分割。

3.1 Heart Detection and Localization

3.2 Heart Segmentation

3.3 Edge-loss Head and Loss Function

由于嵌入式三维U-net的简化，我们通过增加一个边缘损失模块来补偿精度损失。结果表明，边缘损失模块不仅有助于获得更好的分割结果，而且加快了收敛速度。具体地说，我们将固定的3D Sobel核应用在最终的预测图和one-hot真值标签上，计算它们之间的L2损失，称为边缘损失。边缘损失鼓励3D U-net更多地关注心脏器官和背景之间的边缘，这通常更难学习。我们设计了6个对应于6个方向（正x，y，z和负x，y，z）的3D Sobel核，每个核可以描述为3×3×3矩阵，其中中间的3×3矩阵为零,如图4 $S_{middle}$所示的值，上面的3×3矩阵为$S_{up}$所示的值，下面的3×3矩阵为 $S_{down}$所示的值。边缘提取可以通过Sobel核与分割的预测图和真值标签之间的卷积运算来实现。