《多分辨率VB-nets分割CT胸腔器官》论文笔记

原论文：M. Han, G. Yao, W. Zhang, G. Mu, Y. Zhan, X. Zhou, Y. Gao. Segmentation of CT thoracic organs by multi-resolution VB-nets. ISBI 2019.

 

1 简介

在肺癌和食道癌中，放疗是一种常见的治疗手段。在制定放疗方案时，需要把靶区肿瘤和周围的健康器官即危及器官（OARs）严格的分割界定开来。通常都是人工标记轮廓线来完成的。这一步骤通常都要耗费肿瘤医生数小时的时间，并且具有很大的操作上的不确定性。对于某些器官（比如食管），这种分割尤其具有挑战性：对于不同病人这些器官的形状和位置往往不一样；CT影像的轮廓对比度低；有时甚至有缺失的情况。

近年来，基于深度学习的方法被广泛应用于医学图像分割领域。其中以U-Net和V-Net又是最为有名。V-Net结合了U-Net和残差网络的特点，通过更加平滑的梯度流，更加易于优化和收敛。文中提出了一种改进的V-Net模型，命名为VB-Net，用来对靶区肿瘤和OARs进行分割。经过在一些国际上通用的数据集和ISBI SegTHOR 2019挑战赛上的数据集进行验证，文中提出的VB-Net都显示出了很高的准确率、速度和鲁棒性。

 

2 方法

2.1 用于精确器官分割的VB-Net

V-Net的最初提出是用来在MRI图像上利用端到端的全卷积网络来分割前列腺的。它由两个路径组成，左边的收缩路径通过卷积和降采样来提取高层次（high-level）的上下文信息。右边的扩张通道使用跳跃连接（skip connection）将高层次的上下文信息和精细的局部信息结合起来，从而使得边缘定位更加精确。依靠引入残差函数和跳跃连接，V-Net在图像分割的精确度上超越了很多经典的CNN网络。

文中提出的VB-Net的结构如Fig.1所示。将传统的卷积层替换成Fig.1底部所示的瓶颈状结构。基于这种结构特点，命名为VB-Net（B代表bottleneck的首字母缩写）。这种瓶颈状结构包含三个卷积层。第一个卷积层采用1×1×1的卷积核来减少特征图谱的通道数。第二层卷积层主要做时域上的卷积，其卷积核的大小和传统的卷积层一致。最后一个卷积层用1×1×1的卷积核来增加特征图谱的通道数使其恢复到原来的大小。通过对减少通道数之后的特征图谱做时域上的卷积有两个好处：1）大大的减少了模型的大小，举例来说，从V-Net（250MB）减小到VB-Net（8.8MB）；2）模型推断的时间也缩短了。接下来，文中对传统卷积层和瓶颈状结构的模型大小进行了理论分析。一个卷积核大小为K×K×K、输入/输出通道大小为C的卷积层模型大小（参数大小）为 。作为对比，相应的瓶颈状结构的模型大小为，其中N是原始通道大小和缩减之后通道大小的比率。大多数情况下给定K=3，则使用瓶颈状结构后的模型压缩比为 ，当N=2时模型压缩率为3.48，当N=4时模型压缩率则达到了12.34。在下面的实验中，设N等于4，则模型的大小从V-Net的250MB被压缩到8.8MB。当模型的大小缩减后，更方便在云端和移动端部署分割网络。

 

2.2 多分辨率策略

许多深度学习算法都是以一种分辨率来分割器官。由于3D影像（CT、MR）尺寸上都比较大，比如说512×512×300，如果让整个3D图像体积作为网络的输入无疑要占用很多的GPU内存，增加了由于GPU显存不足而报错的可能性。一种解决办法就是对图像体积进行重新采样，降低分辨率，然而这样做会破坏图像的细节，同时分割的边界线也会变成锯齿状。另一种常用的策略是用滑动窗口把图像体积分成一些部分重叠的子序列，然而这种方法会比较耗费时间，因而在工业上并不实用。

文中提出了一种多分辨率策略。针对不同的图像分辨率，单独训练两个VB-net。在低分辨率的情况下，训练一个VB-Net来粗略定位每个器官的感兴趣体积（Volume of interest, VOI）。使用经过重新采样的以6mm大小为一体素（voxel）的图像来训练VB-Net。在高分辨率的情况下，利用检测到的VOI训练VB-Net来得到更精准的器官边界线。

 

 

 

 

2.3 单类/多类损失函数

目前流行的分割方法中有很多种损失函数，比如单像素级交叉熵，focal loss以及Dice loss等等。Dice loss相比交叉熵有一个明显的好处：它更重点关注前景的体素（voxels），而不管整个图像有多少个背景的体素。文中使用了单类/多类均适用的Dice loss函数的一种通用形式，其数学公式如下：

 

 

其中C代表分类标签的数目， 是体素i被分成第c类的概率， 是一个二值标签，表示体素i是否被分为第c类。当C=1时，上式便是Dice损失函数的单类形式。

 

3 数据和实验

3.1 数据集

实验数据来自SegTHOR 2019训练和测试集。训练和测试集分别包括40和20位病人的数据。CT扫描在平面上大小为512×512个像素，在空间上分辨率在0.90mm和1.37mm之间变化。切片数目在150到284之间变化，一个切片的厚度则在2mm至3.7mm之间。最常见的空间上的分辨率是0.98×0.98×2.5。

 

3.2 灰度归一化

为了加速神经网络训练时的收敛，需要首先对图像的灰度进行归一化。由于需要分割的有4个胸腔器官（心脏、主动脉、气管、食管），文中选择纵隔窗来做全局灰度归一化。比如说，假定窗位40，窗宽350，设定灰度值区间分别为-310和400。把图像在这个区间内的灰度值线性归一化到[-1, 1]之间，灰度值小于下限则设为-1，大于上限则设为1。

 

3.3 Patch级网络训练

训练图像首先进行各向同性的重采样，再进行灰度归一化。对于低分辨率的情形，利用6mm的各向同性的空间体对其进行重采样。为了防止标签的丢失，在采样之前，对原图像空间气管和食管的mask进行10mm的扩张。在高分辨率的情况下，利用1mm的空间体进行重采样，并不做任何mask的扩张。经过重采样之后，随机得到的大小为96×96×96体素的3D子图像volume作为训练集。在低分辨率情形时，对整个图像域上随机采样，而在高分辨率的情形下，仅对ground truth所mask的区域进行采样。用这种方式处理，则高分辨率网络更加侧重于器官的边界。对于每个采样的图像crop，相应的mask crop则被作为ground truth的mask来提取，用来作为网络的预测目标。文中分别对低分辨率网络和高分辨率网络单独进行训练。

 

3.4 全卷积网络推理（inference）

在推理（inference）阶段，前面已经提到了一种连接低分辨率和高分辨率网络的策略。低分辨率网络的目的是粗略的分割器官，用来估算VOI。之后，根据得到的VOI来重新采样得到一个高分辨率的区域，然后利用高分辨率的网络来精确分割得到器官的边界。

和一般的patch级网络的训练不同，文中是把整个图像volume作为全卷积网络的输入，而不是像通常那样裁切成若干部分重叠的子volume。之所以具备这样的处理能力，主要是有两点：1）应用了多分辨率策略。低分辨率网络进行了6mm空间体的重采样之后，只占用几百MB的GPU显存。而高分辨率网络又只针对低分辨率网络得到的整个图像的部分区域进行处理。2）除此之外，另一个原因是作者团队使用了scratch来实现推断引擎，而不是使用pytorch、tensorflow这样的开源框架。这样做的好处是可以针对VB-Net来自主对GPU显存使用进行优化，相对于使用pytorch实现的版本，显存占用只是后者的75%。显存优化的细节不在文中探讨范围之内。

 

3.5 后处理

经过高/低分辨率网络的分割之后，通过提取最大3D连通区域来去除噪声区域。对于食管分割，选取大小大于500体素的连通区域。由于其管状结构，因此后处理阶段要特别注意分割中可能出现的连通中断。对于心脏分割，2D切片中可能会有一些小的零散区域，要消除它们，选取切片中的2D最大连通区域即可。