这篇论文是飘在 Arxiv 前十的一篇论文,当时就加入了待看清单,觉得这个自监督学习的方法用于心脏分割挺新颖的。
都知道深度学习需要大量的数据,但是在医学领域获取大量的图像资源是一件很耗费的事情。所以针对数据规模的问题,非监督学习,弱监督学习和自监督学习等方法开始变得热门。这篇论文就是应用自监督学习通过预测解剖位置来完成心脏图像分割。
自我监督学习,它为特征学习制定了一个基于未注释数据的 pretext task。而解剖位置作为一个监督信号,不需要额外的手工注释。
Methods
1. Cardiac MR view plane
一个标准的心脏 MR 扫描是由不同角度的心脏图像组成的,包括短轴、长轴 2 室(2Ch,垂直长轴)、4 室(4Ch,水平长轴)和 3 室(3Ch) 视图。
这些视图被用来评估不同的心脏解剖区域。举例来说,短轴视图显示左心室 (LV) 和右心室 (RV) 的横切面。长轴视图显示心室的间隔和侧壁,以及心房,包括左心房 (LA) 和右心房 (RA)。上图说明了短轴和长轴 2Ch、4Ch 视图是如何面向心脏的。
以往大多数心脏 MR 图像分割的研究都是将短轴和长轴图像分开考虑,忽略了不同图像的相对方向。但在这篇论文中,作者提出利用短轴视图和长轴视图的相对方向,以及由视图平面定义的解剖位置,可以制定一个 pretext task,以自我监督的方式训练网络,提高数据效率。
2. Self-Supervised Learning (SSL)
如上图中的 (b) 所示,2Ch 视图平分 LV 和 RV,而 4Ch 视图平分 LV。它们在 LV 相交。沿着腔室视图线,定义了 9 个解剖位置,由边界框表示,包括交叉点、左边两个框、右边两个框、前面两个框和后面两个框。从左到右和从后到前的方向可以从 DICOM 头文件中获得。pretext task 是预测这九个边界框定义的解剖位置。直觉上来说,要想让神经网络识别出这些解剖位置,它不仅需要了解左右心室的大致位置,还需要了解它们的邻近区域。然而这些学习到的特征可以转移到一个相关的,但要求更高的任务,即准确分割心室任务。
同样,对于长轴 4Ch 视图图像,可以将 2Ch 视图和中短轴视图覆盖在其上,如上图 (c) 所示。沿着 2Ch 视图和中短轴视图线,定义了 9 个解剖位置,包括交点、左侧两个框、右侧两个框、上方两个框和下方两个框。用于长轴图像分析的 pretext task 就是预测这些解剖位置。
为了学习 pretext task,作者训练了一个 10 路分割网络,它将 9 个边界框和背景分割开来。使用标准的 U-net 结构,它由编码器、解码器、skip connection (头尾连接)组成,如下图中的 (a) 所示。并且采用交叉熵作为损失函数。
3. Transfer Learning
网络在 pretext task (task 1) 上进行自训练后,将会转移到新的任务 (task 2) 上,对 LV 腔、心肌腔、RV 腔等解剖结构进行精确分割。为了实现这一点,可以简单地将 task 1 的头部替换为 task 2 的新头部。这里任务头部指的是 U-net 的最后一个卷积层,其与 K 通道输出进行 1×1 卷积,K 表示类别数。
下面是论文中 3 种不同的迁移学习方法:
- SSL Decoder:该方法将任务头切换成 task 2,并冻结 Encoder 中之前学习到的 weights,仅仅对 Decoder 进行微调。
- SSL All:该方法将任务头切换成 task 2,对 Encoder,Decoder 都进行微调。
- SSL MultiTask:该方法是对多个任务进行微调学习,即任务头有两个分支 - task 1 和 task 2。这是为了避免在学习新任务时忘记 pretext task。下面是该方法使用的损失函数:
Experiments and Results
1. Short-axis image segmentation
如下表 Table 1 所示,即使冻结了 Encoder(SSL Decoder 方法)的权重,仅仅微调 Decoder,也能取得比较好的结果来与从头训练的 U-net 进行比较。这表明 SSL 能够在 Encoder 中学习到适合于分割任务的良好特征。
该表也显示出微调所有权重(SSL All 和 SSL Multitask 方法)的方法,其分割精度基本优于从头训练的 U-net,尤其是训练对象数量较少时。平均而言,SSL MultiTask 表现得最好。
下图 Fig.3 是当训练样本逐渐增加时的一些训练结果,可以看到,当 n = 1 时,因为训练样本太少,U-net-scratch(从头训练的 U-net)几乎分割不了图像。相对而言,SSL Multitask 则仍可以分割图像中的大部分区块,即可以分割出 LV(左心房) 和心肌。
当 n 增加到 5 或 10 时,SSL Multitask 在细节方面的性能优于 U-net-scratch,例如,没有 RV 的欠分割错误。然而,当 n 增加到 50 或 100 时,这两种方法的性能相似。下图 Fig.4 为两种方法的定量度量,包括 Dice 度量和每种解剖结构的平均轮廓距离误差。在一个小的数据设置 (n≤10) 下,SSL MultiTask 在所有结构上的性能都优于U-net-scratch。当训练数据越多 (n≥50) 时,两者的性能越接近。
2. Long-axis image segmentation
对长轴图像分割进行了类似的实验。Table 2 展示了 Dice 重叠度量的平均值。结果表明,与 U-net-scratch 相比,SSL 在大多数情况下都能提高分割精度。Fig.5 可视化了分割结果。它表明,在训练数据有限的情况下,SSL MultiTask 通常比 U-net-scratch产生更好的分段。这结果与短轴分割的结果类似。