扫码关注下方公众号:"Python编程与深度学习",领取配套学习资源,并有不定时深度学习相关文章及代码分享。
今天分享一篇发表在CVPR2020上的关于医学图像处理的论文:Deep Distance Transform for Tubular Structure Segmentation in CT Scans (原文链接:[1])。
1 研究背景
医学图像中的管状结构分割(如CT扫描中的血管分割)是使用计算机辅助早期筛查相关疾病的重要步骤。但是目前CT扫描中管状结构的自动分割由于存在对比度差、噪声大、背景复杂等问题而仍然存在很大的挑战。同时,如下图(Figure 1)所示,管状结构其实可以由一系列圆心和半径不断变化的球体组成的。受此启发,这篇文章尝试将这一几何特点融入到管状结构的分割任务中,以提高其分割结果的准确性。
这篇文章主要有如下贡献:
- 提出了一种几何感知的管状结构分割方法——深度距离变换(Deep Distance Transform, DDT),其融合了骨架化的经典距离变换(the classical distance transform for skeletonization)方法和现代深度分割网络
- 在6个医学影像数据集上测试了DDT的性能
2 方法
2.1 整体流程
如上图(Figure 2)所示,DDT在训练阶段和测试阶段有些许差异。
训练阶段:如上图绿色框所示,DDT的输入是3D的CT扫描图$X$,产生两个预测结果:分割概率图$P$和半径分类图$hat{Z}$,并且分别根据标签$Y$和$Z$进行监督学习。
测试阶段: 如上图蓝色框所示,DDT的输入是3D的CT扫描图$X$,同样产生$P$和$hat{Z}$两个预测结果,之后通过$P$中产生伪骨架$S$,再结合$hat{Z}$上对应位置的预测值作为半径,解析出最终的预测结果$hat{Y}$
上述分析显示,DDT的训练需要有两个标签数据:分割图标签$Y$和半径分类图标签$Z$,后面两小节将对标签的构造进行详细介绍。
2.2 分割图标签$hat{Y}$的构造
输入数据为3D的CT扫描图$X$ ($L imes W imes H$),其对应的坐标集合为$V={v|vin N_L imes N_W imes N_H}$。定义$X$上位置$v$的值为$x_v=X(v)$,则$Y$中落在管状结构中的位置$v_{in}$置为$y_{v_{in}}=1$,其余位置$v_{others}$置为$y_{v_{others}}=0$,因此有$y_vin{0,1}$,是一个二类分割任务。
2.3 半径分类图标签$hat{Z}$的构造
上述的分割图标签构造和其它的3D分割方法是一致的,这篇文章的创新在于融入了半径分类图的预测,其作用是为每一个$y_v=1$的位置$v$预测一个半径,从而可以根据Figure 1的几何结构产生分割结果,再和模型预测的分割图综合起来得到最终的预测结果。
以位置$v$为中心的球体对应的半径定义为$v$到最近的分割表面的距离$d_v$。先定义分割表面:
$$C_v={v|y_v=1,exists uin mathcal{N}(v),y_u=0}$$
其中$mathcal{N}(v)$表示位置$v$的6个邻居位置(3D空间中1个位置有6个邻居),如果邻居位置中存在$y_u=0$的,则说明该位置$v$在分割表面上。则$d_v$可以定义为:
$$d_v=left{egin{matrix}min_{uin C_v}left |v-u
ight |_2,if y_v=1\ 0, if y_v=0end{matrix}
ight.$$
上述公式表示:对于分割区域内的每一个位置$v$,计算其与分割表面$C_v$中距离最近的点作为其在半径分割图上$Z(v)$的值,对于分割区域外的位置没有半径表示,$Z(v)=0$。
但是注意到,这样计算出来$d_v$是连续的实数,而[3]中提到,医学影像数据的标注中普遍存在标注误差,如果将$d_v$的学习视为回归问题来训练网络,将由于标注存在的误差而导致网络难以收敛或者产生不稳定的预测结果,如[4]中所提到的。因此文中将$d_v$取整得到离散的数值$z_vin {0,cdots,K}$,从而将半径分类图标签$Z$的学习是视为一个$K$分类问题,即有$Z(v)=z_v$,使得模型更容易训练也具有更好的性能。
2.4 损失函数
如2.1小节中所介绍的,DDT的预测有两部分:分割概率图$P$和半径分类概率图$hat{Z}$,因此在训练中需要分别对这两部分计算损失值,再通过反向传播进行训练。
对于分割图部分,采用加权交叉熵损失函数 (weighted cross-entropy loss function)进行计算:
$$mathcal{L}_{cls}=-sum_{vin V}left(eta_py_vlog p_v(W, w_{cls})+eta_n(1-y_v)log (1-p_v(W, w_{cls}))
ight )$$
其中$W$表示骨干网络 (network backbone)的参数,$w_{cls}$表示分割图预测分支的参数,$p_v(W,w_{cls})$表示位置$v$被预测为管状结构体素的概率值。此外有$eta_p=frac{0.5}{sum_v y_v},eta_n=frac{0.5}{sum_v(1-y_v)}.$
对于半径分类图部分,其损失函数表示为:
$$mathcal{L}_{dis}=-eta_psum_{vin V}sum^K_{k=1}left (mathbf{1}(z_v=k)left (log g_v^k(W,w_{dis})+lambdaomega_vlog (1-max_lg_v^l(W,w_{dis}))
ight )
ight )$$
其中$W$表示骨干网络 (network backbone)的参数,$w_{dis}$表示半径分类图预测分支的参数。$mathbf{1}(cdot)$是指示函数,当$z_v=k$的时候返回1,否则返回0。$lambda$是平衡两个函数项的参数,在文中直接设置$lambda=1$,$g_v^k(W,w_{dis})$表示位置$v$的半径被预测为属于第$k$个类别的概率,$omega_v$是标准化权重,定义为$omega_v=frac{|arg max_lg_v^l(W,w_{dis})-z_v|}{K}$
最终的损失函数表示为:
$$mathcal{L}=mathcal{L}_{cls}+mathcal{L}_{dis}$$
训练的目标表示为:
$$(W^*,w_{cls}^*,w_{dis}^*)=arg min_{W,w_{cls},w_{dis}}mathcal{L}$$
2.5 几何学层面的精调 (Geometry-aware Refinement)
受[5]的启发,文中提出在几何学层面的精调方法 (gometry-aware refinement, GAR),具体步骤如下:
1. 伪骨架生成 (Pseudo skeleton generation):预测分割图$P$通过阈值$T^p$过滤出伪骨架$S$,具体地有:$p_v>T^p,s_v=1;otherwise,s_v=0$
2. 形状重建 (Shape reconstruction):如Figure 1所示,对于每一个位置$v$,其对应的半径为$hat{z}_v=arg max_kg_v^k$,文中通过高斯分布生成对应的球体,多个球体合起来即可得到软化后的结构形状$ ilde{Y}^s$ (soft reconstructed shape):$ ilde{y}_v^s=sum_{uin{u'|s_{u'}>0}}c_uPhi(v;u,Sigma_u)$,其中$Phi(cdot)$表示多元高斯分布的密度函数,$u$表示均值,$Sigma_u=left (frac{hat{z}_u}{3}
ight )^2I$表示协方差矩阵。此外,采用标准化因子$c_u=sqrt{(2pi)^3det(Sigma_u)}$放缩$Phi(cdot)$的结果。这一步可以得到基于$P$和$hat{Z}$得到的形状重建结果$ ilde{Y}$
3. 精调分割结果 (Segmentation refinement):对于模型预测的分割概率图$P$,可以采用上一步得到的结果$ ilde{Y}$进行调整:$ ilde{y}_v^r=sum_{uin {u'|s_{u'}>0}}p_u c_uPhi(v;u,Sigma_u)$,然后采用阈值过滤$ ilde{y}^r_v>T^r,hat{y}_v=1$,$otherwise,hat{y}_v=0$得到最终预测结果$hat{Y}$
3 实验结果
这里我只给出论文中的部分实验结果,具体的实验结果分析以及实验和参数的设置请看原文。
4 参考资料
[1] http://openaccess.thecvf.com/content_CVPR_2020/papers/Wang_Deep_Distance_Transform_for_Tubular_Structure_Segmentation_in_CT_Scans_CVPR_2020_paper.pdf
[2] Fethallah Benmansour and Laurent D. Cohen. Tubular structure segmentation based on minimal path method and anisotropic enhancement. International Journal of Computer Vision, 92(2):192–210, 2011.
[3] Yan Wang, Yuyin Zhou, Peng Tang, Wei Shen, Elliot K. Fishman, and Alan L. Yuille. Training multi-organ segmentation networks with sample selection by relaxed upper confident bound. In Proc. MICCAI, 2018.
[4] Rasmus Rothe, Radu Timofte, and Luc Van Gool. Deep expectation of real and apparent age from a single image without facial landmarks. International Journal of Computer Vision, 126(2-4):144–157, 2018.
[5] Wei Shen, Kai Zhao, Yuan Jiang, Yan Wang, Xiang Bai, and Alan L. Yuille. Deepskeleton: Learning multi-task scale-associated deep side outputs for object skeleton extraction in natural images. IEEE Trans. Image Processing, 26(11):5298–5311, 2017.