【阅读笔记】《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》的笔记，代码。更新于2019.06.24。

摘要

Introduction

心脏CT体元（cardiac CT volumes）中的心脏分割可以应用于多种医疗领域。难点：由于GPU限制，图片被分成小部分进行分割，这种方法求是了全局约束，而且由于心脏只占全局的一小部分，因此也浪费了大量的计算资源。

本文提出了一种基于检测的分割方法，从而让分割网络能够关注CT图像中最相关的部分，而不需要处理整个图像。CFUN的检测和定位网络基于Faster R-CNN，分割网络基于3D U-Net，同时也做了一些重要的修改。

与传统Faster R-CNN处理多个不同目标的boxes不同，这里强迫检测网络只能识别一个包括了整个心脏标注结构的bounding box。考虑到CT图像特性和3D边界损失，Faster R-CNN中的ResNet结构被替换成了P3D Bottleneck。

Related Work

Traditional Methods

3D U-Net Based Segmentation Network

3D U-net直接以3D体元为输入。尽管效果不错，但是其需要结合tiling strategy。具体而言，将3D体元切割成有顺序的小图块，分别送入神经网络，再将输出缝合起来得到最终的分割结果。考虑到心脏图片中的大部分都属于背景，这种操作会浪费大量的计算量。

可以两段处理：第一段分割粗略的心脏图像，第二部分输出原始分辨率下的分割结果。

Faster R-CNN and Region-proposal Based Segmentation

网络可以同时实现检测、分类和分割，也有一些尝试将Mask R-CNN应用到医学图像分割。

Method

CFUN包括连续的两部分：一个修改过的3D Faster R-CNN来生成以心脏为中心的bounding box；和一个3D U-net，用Edge-loss head计算边缘信息并比较其与真实边界来分割以心脏为中心的图片。测试过程不需要Edge-loss head。

Heart Detection and Localization

通过用3D kernel替换所有的2D kernel将Faster R-CNN扩展到3维。由于$z$轴上的分辨率与其他两个轴是不同的，因此用P3D ResNet替换了Faster R-CNN中的ResNet。

同时，在P3D网络后面还增加了FPN（Feature Pyramid Networks）来整合不同分辨率下的特征图。在CFUN网络中，这两个子网络组成了RPN（region proposal network）部分。

Heart Segmentation

将修改后的3D U-net结构作为CFUN网络的分割部分。与Mask R-CNN整合特征图作为分割网络的做法不同，这里级联了在bounding box内的原始的CT图像，这里可以韩做事分割的一个独立部分。尽管Mask R-CNN证明了可以同时学习检测和分割的特征，但是作者的剥离实验证明了考虑到large RoI-align size（64x64x64），这么做的难度大且推理效果不太好。这里借鉴了深度监督（深度残差）的概念。

Edge-loss Head and Loss Function

受到这篇论文的启发，在精度损失上有增加了边缘损失，其证明了这么做不仅能够增加精度，而且有助于加速收敛。具体而言，作者应用了固定的3D Sobel-kernels在最终的估计图上，同时也包括原来的one-hot真值之间的$l_2$损失计算。

共涉及了6个3D Sobel-kernels，分别对应6个方向（$x,y,z$和其反向），每个核尺寸均为$3\times 3\times 3$，中间的$3\times 3$矩阵$S_{middle}$是0，上面的$3\times 3$矩阵$S_{up}$由下图给定，而下面的$3\times 3$矩阵$S_{down}$有对应的上矩阵给定，取相反数。边缘提取可以看成是分割估计图或真值与Sobel kernel的卷积操作。

边缘损失定义如下：

其中，$C$是分割的类别，$K$是Sobel kernel数，$M$是图像体元个数，$y$是one-hot label map，$p$是网络的估计图。$E_k(y{)}_m^c$表示从图像$y$的第$c$个维度提取第$k$个边缘信息后的第$m$个像素值。

CFUN优化四个损失函数：positive bounding box $L_{seg}$的分割损失，head $L_{edge}$的边缘损失，回归损失和$L_{box}$和$L_{cls}$的分类损失。总损失定义为：

其中$w_1,w_2,w_3$和$w_4$是平衡不同损失的权重。在本文实验中，$w_1:w_2:w_3:w_4=2:2:2:1$。

Experiments

Dataset

数据库：MM-WHS2017挑战赛，包括20个训练体元和40个测试体元。包括7个标注：左心室血腔（the left ventricle blood cavity, LV），左心室心肌（the myocardium of the left ventricle, Myo），右心室血腔（the right ventricle blood cavity, RV），左心房血腔（the left atrium blood cavity, LA），右心房血腔（te right atrium blood cavity, RA），升主动脉（the ascending aorta, AA），和肺动脉（the pulmonary artery, PA）。

由于只有20个训练样本，但是网络训练需要更多，所以论文作者对样本进行了重新分配。用40个原始测试样本作为训练样本，5个随机训练样本作为验证集，剩下的15个原始训练样本作为测试集。

Training strategy and detail

架构：PyTorch
设备：NVIDIA GeForce GTX 1080Ti GPU
尺寸：resize成320x320x192，RoI-align尺寸64x64x64
优化：Adam（bz=1，lr=0.001，每个epoch的步长32）

Testing results

分割一个心脏的时间小于15s。

Ablation experiments

Different feature maps as the input of 3D U-net

w/o Edge-loss head and w/o bounding box refine head

Conclusion

代码

R-CNN的定位能力准确，U-Net的分割能力强大，所以CFUN设计了一个结合二者的网络结构，同时实现检测和分割推理获得整个心脏的分割结果，速度和精度俱佳。同时，基于边缘信息设计了一个新的损失函数，称为3D Edge_loss，用于加速收敛并提升分割结果。

网络结构

数据集主要基于MM-WHS2017挑战赛，但也可以用其他数据集。如果要用自己的数据库，需要注意CFUN和Mask-RCNN的区别。在这个流程中，真值bounding box是整个心脏器官，而不是心脏的具体位置，比如左心室血腔（left ventricle blood cavity, LV）、左心室肌（myocardium of the left ventricle, Myo）等，因为其紧密联系。因此，如果数据库是分割不同器官，则需要修改model.py代码中的load_image_gt函数，来生成特定于器官的不同bounding-box。

首先，需要准备一个dataset.json文件，其中包括训练和测试图像的所有方向以及训练标签。