3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation 论文解读与程序复现

时间 2021-01-11 标签医学影像分割 Unet 语义分割人工智能机器学习

转载请注明出处：http://www.noobyard.com/article/p-bcjkbhwr-sy.html

这篇论文题为3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation。

论文解读

在生物医学领域，3D数据是很多的，一层一层转化为2D数据去标注训是不现实的，而且用整个3D体积的全部数据去训练既低效又极有可能过拟合。这篇文章提出的3D Unet只需要少部分2D的标注切片就可以生成密集的立体的分割。此网络主要有两个不同的作用方式，一是在一个稀疏标注的数据集上训练并在此数据集上预测其他未标注的地方，另一个方法是在多个稀疏标注的数据集训练，然后预测新的数据。

3D-Unet的结构基本上和2D一模一样（具体结构见我的上一篇博客），只是增加了一个维度。值得注意的事，这里作者还用了Batch Normalization 防止梯度爆炸，并且在BN后增加了缩放和平移： $x_{new}=\alpha\cdot x+\beta$ ，其中两个超参是学习出来的。

本文作者用的训练数据是爪蟾肾胚，标记时0表示在管内，1表示管壁，2表示背景，3表示未标记。训练时，背景的softmax损失表达式中的权重被减少，而管内的权重增加，这样做有利于更好地分割，因为绝大多数地方是背景，相当于起到了突出强调作用。
在实验过程中，作者还采取了交叉验证和对比验证，对半自动和全自动分割两种情况，最好的是3D+BN，其次是3D（无BN），最后是2D+BN。

程序复现

参考代码：https://github.com/ellisdg/3DUnetCNN
我用的数据集是BRATS2017，是用来分割脑部肿瘤的，详见http://www.med.upenn.edu/sbia/brats2017.html
我使用的是改良后的model：Isensee et al. 2017, 配置NVIDIA GTX TITAN X(12 G)训练时长约一天，代码中一共500个epoch，训练到第256个完成训练：

loss随时间变化：

测试结果：Dice 系数随分布图：

图中Whole Tumor是测试中所有标注区域，Tumor core是标签1 和4的区域，enchancing tumor是标签4的区域。这与Brats 竞赛的评分机制是一样的。Dice系数在我的理解里类似于IoU，体现分割精确度，其定义式为 $DC=\frac{2|X\cap Y|}{|X|+|Y|}$

从图中可见效果还是很好的。

此外我还对一些切片进行了可视化，显示预测效果。prediction保存形是.nii 文件，需要用MATLAB 安装Tools for NIfTI (ANALYZE) MR image包进行可视化。我选取了一个样本的16个切片进行预测结果和ground truth对比：

图中前两行是prediction，后两行是groundtruth。可以看出预测结果还是十分准确的。