適用於精細結構的語義分割網絡
摘要
電子顯微鏡(EM)是一種廣泛使用的技術,用於獲取神經組織的高分辨率圖像,以可視化神經元結構,如化學突觸和囊泡。爲了獲得組織更精細的結構,超高分辨率的電鏡圖像已經越來越常見,但如何對這些圖像進行有效的分割卻是一個挑戰。由於顯存的限制,直接將原始圖像輸入網絡進行訓練是不允許的,而將圖像裁剪成一堆小塊進行訓練又會損失圖像的全局信息。如果直接將圖像進行下采樣然後再進行訓練,這樣又會損失圖像的細節信息。如何在超高分辨圖像上同時結合全局信息和細節信息是一個挑戰。本文主要提出多尺度進行超高分辨率圖像分割的方法,分別在三種尺度上(原始分辨率,下采樣二倍和下采樣四倍)獨立進行網絡訓練,最後再進行不同分辨率下分割結果的融合。與單獨的裁剪方法和下采樣方法相比,我們提出的方法在保持全局信息的同時又能保持細節信息。另一方面,我們採用了兩種不同模型的融合,每種模型使用不同的損失函數以達到對神經元邊界分割的魯棒性。最終實驗在complex賽道上(https://www.biendata.com/competition/urisc/)獲得第四名,並公佈了實驗代碼:https://github.com/weihuang527/neuron_challenge。
關鍵詞:神經元分割、 神經網絡、 多尺度、融合
數據處理
訓練數據準備
主辦方提供20張大小爲9958x9959帶有手工標註結果(groundtruth)的訓練數據集,還有10張同樣大小不帶有groundtruth的測試集。在實驗過程中,我們在20張訓練集中挑選4張作爲驗證集,剩餘的16張作爲訓練集。
因爲主辦方提供的圖像大小爲9958x9959,不能被1024整除,爲了後續實驗的方便,我們統一將所有圖像用0擴充成10240x10240的大小。
爲了不同尺度網絡可以同時進行訓練,我們將擴充後的數據在三種尺度上各準備一份,包含原始分辨率、下采樣2倍和下采樣4倍。在三種分辨率上,我們分別把圖像裁剪成大小爲1024x1024,重疊512的圖像塊。
訓練數據集增強
爲了充分擴充訓練數據集,在網絡訓練階段,網絡每次迭代過程中隨機從裁剪塊中選擇一張,然後將1024x1024的圖像隨機裁剪成512x512的圖像送進網絡,這樣可以充分保證訓練數據的多樣性。
當然本實驗還包含一些常規的訓練數據增強的方法,包含隨機翻轉(水平和豎直兩個方向),隨機旋轉(90度、180度和270度三個方向),隨機尺度變換(100像素內),隨機顏色變換(對比度範圍0.2,亮度範圍0.2,飽和度範圍0.2)。如圖1所示,顯示了三種分辨率下經過上述增強方法後的訓練圖像和對應的groundtruth。
實驗方法
實驗結構
爲了得到二值化的結果,我們取閾值0.65,將小於等於0.65的像素值變爲0(黑色邊界),大於0.65的像素值變爲1(白色背景)。
Fusionnet網絡
如圖3所示,Fusionnet與傳統的Unet相比,它不僅包含了一個收縮路徑以提取特徵和一個對稱的擴張路徑以更好地定位,還在每層下采樣和上採樣之間嵌入了殘差模塊,同時了保留了Unet的跳層連接。從文獻[1]中可以看出,Fusionnet比Unet在神經元邊界提取方面,有着更強的優勢。所以本文采取了文獻[1]的fusionnet結構,網絡第一個殘差塊的輸出特徵通道數爲32,後面每下采樣一次特徵通道數增加2倍,共有4層下采樣層。相同也有4層上採樣層,每上採樣一次特徵通道數縮小2倍,網絡的最後一層爲Sigmoid函數,確保最後輸出的概率圖範圍在0和1之間。
和文獻[1]一樣,我們的損失函數採用單像素的MSE函數,該損失函數能保證輸出邊界的寬度,不會比groundtruth更粗或者更細。
該模型在4張NVIDIA GTX XP上進行訓練,在驗證集上選擇最優模型(原始分辨下大約需要迭代100000次,下采樣2倍下與大約需要迭代40000次,下采樣4倍下大約需要迭代35000次)。Batch size爲16,學習率爲0.0005,採用學習率逐漸衰減,衰減下限爲0.00001。優化器採用Adam模型,beta1=0.9,beta2=0.999。三種不同分辨率下網絡的結構以及參數保持一致沒有任何修改。
Unet++網絡
如圖4所示,我們選擇Unet的另一個變形-Unet++[2]作爲我們的分割模型。它與標準的Unet相比,它具有很多的跳層連接和不同尺度下的密集結合,極大地提升了網絡提取特徵的能力,並不同尺度下的訓練提升了模型對不同尺度的魯棒性,非常適用於神經元圖像的分析,因爲不同的神經元具有不用的大小,而多尺度訓練正是克服了這一難點。同時我們將裏面的BatchNorm換成了InstanceNorm,採用了Res-Net50作爲backbone,因爲發現Res-Net101會發生過擬合現象,Res-Net34效果不如Res-Net50。
我們採用了Cross Entropy(CE)和Dice兩種函數的加權作爲損失函數,其中CE的權重爲0.1,Dice的權重爲1。考慮到神經元邊界提取中類別不平衡的問題(邊界像素僅佔整張圖像的約5%),我們在CE中添加了類別的權重,權重值爲各自的在圖像中的佔比。而Dice因爲自帶這種類別的平衡,所以我們並沒有在Dice上添加權重。
同樣該模型在4張NVIDIA GTX XP上進行訓練,在驗證集上選擇最優模型(原始分辨下大約需要迭代130000次,下采樣2倍下與大約需要迭代150000次,下采樣4倍下大約需要迭代30000次)。Batch size爲8,學習率爲0.01,沒有采用學習率逐漸衰減。優化器採用Adam模型,beta1=0.9,beta2=0.999。三種不同分辨率下網絡的結構以及參數保持一致沒有任何修改。
實驗細節
本實驗採用2折交叉驗證,即另外從10張訓練集中選擇4張出來作爲驗證集,然後進行上述同樣的訓練,在測試階段,將第1折的結果和第2折的結果直接取平均求和,然後再進行最後的加權求和。
在測試階段,我們也採用了測試集增強方案,由於時間的限制,我們在下采樣2倍和下采樣4倍的模型上採用旋轉90度、180度和270度,三種方式的增強,而在原始分辨率上不進行增強。
爲了節省預測時間,在下采樣4倍的模型下,預測時輸入的圖像大小爲2048x2048,在下采樣2倍和原始分辨率下,預測時輸入的圖像大小爲1024x1024。
實驗結果
定性分析
如圖5所示,同樣的模型在不同的分辨率下也有着不同的分割結果,從預測結果中,我們可以看出在低分辨率下,模型對全局信息把握得更好,細胞內的誤判的邊界更少或者概率更低,在全分辨率下,細節信息更準確,體現在邊界粗細更準確,但缺少全局信息,導致誤判的細胞內部邊界更多。
定量分析
多尺度融合的有效性
模型 | 閾值 | 測試分數 |
---|---|---|
Fusionnet_x0 | 0.7 | 0.63870 |
Fusionnet_x2 | 0.7 | 0.62883 |
Fusionnet_x4 | 0.7 | 0.59674 |
W_F_X0 | W_F_X2 | W_F_X4 | 閾值 | 測試分數 |
---|---|---|---|---|
0.33 | 0.33 | 0.33 | 0.7 | 0.64301 |
0.50 | 0.30 | 0.20 | 0.7 | 0.64566 |
0.40 | 0.35 | 0.25 | 0.7 | 0.64490 |
從表1、2中我們可以發現,三種尺度下的融合可以極大的提升Fusionnet的分割結果。(我們所有的評測指標都在F1下進行)
多模型融合的有效性
模型 | Fusionnet | Unet++ | 兩者融合 |
---|---|---|---|
測試分數 | 0.39089 | 0.44362 | 0.45007 |
從表3中我們可以看出,兩種不同的模型的融合可以顯著提升分割分數。
權重以及閾值選擇
W_F_X0 | W_F_X2 | W_F_X4 | W_U_X0 | W_U_X2 | W_U_X4 | 閾值 | 分數 |
---|---|---|---|---|---|---|---|
0.17 | 0.17 | 0.17 | 0.17 | 0.17 | 0.17 | 0.70 | 0.64677 |
0.17 | 0.17 | 0.17 | 0.17 | 0.17 | 0.17 | 0.60 | 0.64506 |
0.17 | 0.17 | 0.17 | 0.17 | 0.17 | 0.17 | 0.65 | 0.64955 |
0.25 | 0.15 | 0.10 | 0.10 | 0.15 | 0.25 | 0.65 | 0.65033 |
0.25 | 0.15 | 0.10 | 0.10 | 0.25 | 0.15 | 0.65 | 0.65116 |
從表4中我們可以發現,當閾值取0.65和權重分別爲0.25、0.15、0.10、0.10、0.25、0.10時可以在測試集上取得相對較好的分割分數。
結論
優點:本實驗考慮不同尺度和不同模型的融合,結合神經元圖像的全局信息和局部細節信息,針對不同大小的神經元有着更好地魯棒性,同時在兩種不同的網絡中使用不同的損失函數,可以彌補每個損失函數的缺陷,從而達到取長補短的效果。
缺點:我們發現不同的閾值對着最後的分割結果有着顯著的影響,而我們對閾值的選擇有一定的盲目性,只能憑藉經驗去選擇。另一方面,六個模型求和的權重也是一個超參,目前也只能憑藉進行選擇。
參考文獻
[1] Quan T M, Hildebrand D G C, Jeong W K. Fusionnet: A deep fully residual convolutional neural network for image segmentation in connectomics[J]. arXiv preprint arXiv:1612.05360, 2016. https://arxiv.org/abs/1612.05360
[2] Zhou Z, Siddiquee M M R, Tajbakhsh N, et al. Unet++: A nested u-net architecture for medical image segmentation[M]//Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support. Springer, Cham, 2018: 3-11.
[3] Fakhry A, Zeng T, Ji S. Residual deconvolutional networks for brain electron microscopy image segmentation[J]. IEEE transactions on medical imaging, 2016, 36(2): 447-456. https://ieeexplore.ieee.org/abstract/document/7575638