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。可以看出預測結果還是十分準確的。