基于EM 路由算法的醫學圖像分割UCaps 網絡

王文欣，賀煜航，陳 剛

（武漢大學國家網絡安全學院，武漢 430072）

0 概述

圖像分割是指將圖像分成若干具有相似性質區域的過程。傳統的分割方法包括閾值分割、邊緣檢測、聚類、遺傳算法、區域生長分割法、深度學習等［1］。醫學影像分割［2-3］作為醫學影像分析的核心任務，需要有經驗的醫生對臟器、病灶進行精準分割。當圖像存在多個目標物體時，深度學習分割多個目標就變成了語義解析的任務。在當前醫學影像分割中，語義目標識別的主要任務是對主要骨結構（肋骨、椎骨、頭骨等）、主要器官（肺部、心臟、肝臟等）以及一些解剖標記點（如肝臟頂部、主動脈弓）進行識別、配準和分割，采用深度學習進行醫學影像分割時，需要區分病灶或器官。但是醫學影像分割所需較高的分割精度和穩定性，不正確或不穩定的分割將會直接影響后續計算，從而失去分割意義。

全卷積網 絡（Fully Convolutional Network，FCN）［4］通過將深層網絡的全連接層改為卷積層，從而保存分割的位置信息，分割結果取決于池化層輸出的上采樣倍數，分割精度與分割后語義信息呈負相關。文獻［5］提出U-Net 網絡將跨越連接原理應用于分割問題，利用底層信息補充高層信息，有效提高分割精度。在心臟和大血管分割的問題上［6］，稠密連接的卷積網絡（DenseNet）［7］將任意兩層網絡連接起來，以最大程度地保存網絡信息流和梯度流，降低參數量的同時也解決了梯度消失問題。

FCN、U-Net 以及各種分割網絡都基于卷積神經網 絡（Convolutional Neural Network，CNN）［8］，而CNN 本身存在著一定的弊端。CNN 采用卷積-池化方式并不會保存特征相對于整體的空間依賴關系，即存在同變性和不變性問題，CNN 對旋轉是沒有不變性的。不變性是指目標外觀發生平移、旋轉、光照等變化但人們依舊能夠將其識別出來。這種相對的空間依賴關系對醫療圖像識別、分割十分重要，例如，由左肺、右肺與心臟構成的X 射線圖像，這些圖像在顯示上相互重疊，并具有較嚴謹的空間位置關系。CNN 利用池化操作不僅導致空間上相對信息丟失，而且有可能直接丟棄相關的信息。

FCN、U-Net 等分割網絡利用底層信息補充高層信息的方法解決CNN 存在的問題。另一種方法是研究人員對數據進行處理，例如在訓練數據集中加入旋轉的圖像，使網絡進行學習，從而解決CNN 對旋轉、平移等不敏感的問題。傳統CNN 對于單一目標的旋轉、平移（例如一幅圖片中的眼睛與鼻子互換位置）仍會檢測為正常目標。文獻［9］提出的膠囊網絡是一種在最低層的膠囊（即神經元）上包含目標位置、姿態等信息的神經網絡，其能從本質上優化傳統CNN。

本文結合U-Net 網絡與膠囊網絡提出一種多標簽語義分割網絡UCaps，利用高斯混合模型與EM 路由算法聚合低層膠囊對高層膠囊的概率推導過程，從而保留低層膠囊的精細度和各特征間位置、姿態的統一性，解決傳統分割網絡中卷積與池化分割輪廓不清晰的問題。

1 相關工作

1.1 U-Net 網絡

U-Net 網絡［5］特點是U 型的網絡結構與跨越連接操作，通過下采樣pooling 與上采樣反卷積操作得到顯著特征，同時用跨越連接將前層圖與后層圖相疊加，盡可能少地丟失圖像信息，將信息丟失多而特征明顯的層與信息丟失少特征模糊的層疊加，平衡語義信息和圖像精度。由于醫學圖像數據量通常只有幾百幅，因此底層特征尤為重要。過多的訓練會出現過擬合現象，需要盡量少的數據使其盡快收斂。U-Net 網絡結構如圖1 所示。

圖1 U-Net 網絡結構Fig.1 Structure of U-Net network

U-Net 網絡作為醫學圖像分割領域效果較優的網絡結構，模型代表有結合DenseNet 原理的UNet++［10］、具有嵌套 和Dense 跨越連接結構 的UNet3+［11］等。U-Net 網絡被證明在小數據集、高精度的醫學影像分割中具有相對最優的結構，對構建UCaps 網絡結構具有啟發作用。

1.2 膠囊網絡

文獻［9］使用多維向量代替傳統的標量，提出膠囊網絡。與CNN 相比，膠囊網絡通過動態計算權重在每個正向通道上的耦合系數來對每個目標貢獻不同的權重值。這種優化的機制即動態路由。膠囊網絡與動態路由的基本原理是某一層的膠囊網絡通過變換矩陣對更高級膠囊網絡的實例化參數進行預測，當動態路由使多個預測都一致時，更高級膠囊網絡將會活躍。

動態路由是一種聚類的迭代算法，膠囊網絡的輸出是對動態路由輸入進行聚類的結果。整個動態路由的算法可以用其他聚類算法計算達到更優的效果。

文獻［12］通過矩陣與激活值的方式代替向量并作為基礎膠囊網絡，使用高斯混合模型［13］與EM 算法代替原本的動態路由算法，稱為EM 路由算法。EM 路由算法利用聚類方法對膠囊網絡進行分組，從而形成低層與高層之間的關系，通過聚類低層膠囊對高層膠囊進行推測，即使用EM 路由算法將相應投票非常接近的膠囊網絡進行聚類，得到更高一層的膠囊網絡。膠囊網絡是基于加權的歐氏距離構建高斯混合模型，其中歐氏距離進行聚類的特點在于聚類中心向量是類內向量加權平均。研究人員最早提出，利用向量簇構建膠囊網絡，使用向量模長表示特征的顯著程度，因此向量不適用于EM 路由算法。膠囊結構在改為矩陣時，需要一個標量激活值作為衡量特征顯著程度的標準。具有EM 路由的膠囊網絡結構如圖2 所示。

圖2 具有EM 路由的膠囊網絡結構Fig.2 Structure of capsule network with EM-Routing

膠囊網絡多用于解決分類問題。膠囊網絡之間路由的原理是低級膠囊對高級膠囊形態推測的過程，同樣可以應用于解決圖像分割問題。文獻［9］對兩個手寫體數字重疊的數據集進行識別，基于動態路由的膠囊網絡對重疊數字的識別和分割效果較優。膠囊網絡與其路由方法同樣適用于圖像分割領域。研究人員結合醫學影像（例如肺癌篩查［14］、左心室分割［15］）與膠囊網絡在圖像分割領域上取得較優的成果。文獻［16］提出SegCaps 將膠囊網絡應用在二類分割任務中，通過實驗對其改進后發現，SegCaps 網絡架構雖然大幅減少了訓練參數，但在多標簽分割任務上的準確率大幅下降。

因此，本文研究的重點是將膠囊原理和最適合分割的U 型網絡架構相結合，并融合全新的膠囊路由，使得膠囊網絡在對多標簽進行分割時效果較優。

2 UCaps 網絡結構設計

UCaps 使用共計14 層的網絡結構，其中有1 個卷積膠囊層，9 個EM 路由膠囊層，3 個上采樣膠囊層，1 個分類膠囊層。整體網絡結構取自U-Net 網絡［5］。

UCaps 網絡結構如圖3 所示，由下采樣（左側）和上采樣（右側）組成。UCaps 網絡結構的第1 層使用卷積膠囊層，輸入的圖像卷積為包含［4，4］的姿態矩陣和1 個標量激活值的多個膠囊。

圖3 UCaps 網絡結構Fig.3 Structure of UCaps network

UCaps 網絡將膠囊網絡輸入EM 膠囊層，膠囊網絡進行深度可分離卷積，并將展開的膠囊網絡乘視角不變的變換矩陣W得到投票矩陣V，最后通過EM算法［17］計算高層膠囊網絡的輸出姿態矩陣和激活值，并將結果保存在skip 中用于準備后面跨越連接操作。在每個下采樣步驟，UCaps 網絡使用與傳統膠囊網絡［18］相同結構，逐步抽象其特征，使得膠囊網絡獲取更高級的語義信息。下采樣膠囊為3 層，過大的層數會導致計算參數量太大，過少的層數使得模型缺少足夠的表達能力。上采樣膠囊層類似UNet 網絡中上采樣的過程。跨越連接操作在第4 維度將skip 膠囊與輸出膠囊進行拼接，即將Caps 層輸出添加進DeCaps 層的channel 中。這種做法與UNet 網絡上的通道拼接類似，形成更多的膠囊層，缺點在于使用concat 操作比較占用顯存，優點是融合多尺度維度信息。拼接后的膠囊層融合了同尺度的輸出和上一層膠囊表達的信息。這樣的連接貫穿整個網絡，通過3 次拼接使得UCaps 網絡同時保留低層膠囊的特征信息與高層膠囊的信息。上采樣膠囊層與下采樣對應，同樣為3 層。UCaps 網絡經過3 個上采樣層與膠囊層將結果輸入分類膠囊層，分類膠囊層主要按照分割的種類將對應膠囊輸出進行分類操作，輸出每個膠囊所代表的圖片像素對應的分割標簽類別并預測投票值，進而得到整幅圖的分割預測結果。

EM 路由算法使用高斯分布進行聚類，利用高斯模型對高層膠囊的姿態矩陣進行建模。姿態矩陣為［4，4］，即使用16μ和σ的高斯模型。μ用于估算激活膠囊成本，成本越低越有可能激活高級膠囊。

2.1 上采樣膠囊層

本文在膠囊層基礎上設計上采樣算法作為上采樣膠囊層。上采樣膠囊層示意圖如圖4 所示。

圖4 上采樣膠囊層示意圖Fig.4 Schematic diagram of up-sampling capsule layers

上采樣層過程如式（1）所示：

將Rij帶入到EM算法中，最大化M步驟如式（2）所示：

期望E 步驟的計算如式（3）所示：

2.2 損失函數

上采樣膠囊層使用如下公式作為損失函數：

margin 值隨著訓練次數增加而增大（本文實驗取0.2～0.9）。作為最終膠囊激活值，在訓練早期值小，說明通過高斯聚類后的模型預測處于均勻分布狀態。margin 值過大，使Lloss值被margin 值影響過多。隨著訓練輪數的增加，值會越來越大，說明膠囊激活值指向越來越清晰，margin 值隨之增大，以穩定Lloss值。

2.3 EM 路由與高斯混合模型

高斯混合模型將數據點聚類為混合高斯分布，EM（期望最大化）路由的目的是使用聚類技術將膠囊分組形成一個部分-整體關系。文獻［12］指出，高層膠囊表示其上一層低層膠囊的高斯分布期望，通過高斯混合模型擬合其最小化的損失函數。同層間不同的高斯混合模型反映了相同低層膠囊對不同高層膠囊的傾向性，即高斯混合模型集中了高層膠囊中的低層信息。EM 路由通過訓練視角不變的變換矩陣W，使得高層膠囊不僅包含低層膠囊的特征，同時保留低層膠囊間的位置信息，即使圖片發生旋轉，仍然能將相同位置關系的低層膠囊進行聚類。

3 實驗

3.1 數據集與預處理

數據集使用的是日本放射技術學會（Japanese Society of Radiological Technology，JSRT）數據庫［19］。該數據庫由JSRT 與日本放射學會（JRS）合作創建的有無胸肺結節的標準數字圖像構成。數據庫共有154 張有結節與94 張沒有結節的圖像。圖像分辨率為2 048 像素×2 048 像素，像素尺寸為0.175 mm。標簽圖像來自文獻［20］中對JSRT 數據集進行手動標注的肺部、心臟和鎖骨的標簽圖，將圖像壓縮為128 像素×128 像素，并進行數據集拓展。數據集隨機進行角度偏移、水平或豎直移動、縮放，拓展至2 726 張。數據集圖像與標簽如圖5 所示。

圖5 數據集圖像與標簽Fig.5 Image and label of dataset

3.2 實驗環境與對比模型設計

本文使用設備CPU 為4 核32 GB內存，GPU 使用Tesla-V100-SXM-32GB。實驗環境Tensorflow 1.13.1+cuda10.0+cudnn7.4.5。

本文網絡分別與以下基準網絡進行實驗對比：1）U-Net 網絡［5］，使用CNN 卷積操作的U-Net 網絡，線性上采樣過程取代反卷積層，并用交叉熵評估損失；2）SegCaps 網絡［16］，將兩種類別分割的SegCaps分割網絡改進為多標簽分割，并使用weighted softmax 損失函數進行評估；3）MatVec-Caps 網絡［21］，使用MatVec-CapsNet 網絡進行多標簽分割，MatVec-CapsNet 采用雙路由方法；4）UCaps-Dynamic Routing網絡，使用文獻［9］提出的Dynamic Routing 的路由方法設計UCaps 網絡。

本文將UCaps 網絡與以上4 種網絡進行對比，從CNN 與膠囊、多標簽分割與單標簽分割、EM 路由與動態路由3 個角度進行實驗對比，評估UCaps 網絡的性能。

3.3 實驗結果分析

UCaps 在數據訓練后快速收斂，無論訓練集還是測試集損失分布合理，并未出現欠擬合或過擬合現象。UCaps 對左右肺原始圖片中突出的部分進行分割的精度高，左右鎖骨次之，而在心臟部分的分割準確率約為87%。心臟在圖像中的位置相對不固定，其形狀結構差異性較大，難以對其進行精確分割。隨著訓練迭代次數增多，UCaps 在測試集上心臟部分的分割準確度有顯著提高。UCaps 網絡在測試集上分割結果如圖6 所示，圖6（a）和圖6（c）的底色為原始圖像，圖6（b）和圖6（d）的底色為標簽圖。

圖6 UCaps 網絡的分割結果Fig.6 Segmentation results of UCaps network

U-Net、SegCaps、MaVec-Caps、UCaps 網絡的分割準確率對比如表1 所示。從表1 可以看出，所有網絡在心臟分割部分上的識別率均低于90%。這可能與訓練時間不足有關。本文實驗的目的是在有限的數據集和迭代的情況下獲得更快的收斂速度和更高的準確率。由于心臟在這5 個分類標簽中變化最大、形狀最多，因此對本文的實驗結果造成了一定的影響。SegCaps 網絡主要解決二分類分割問題，本文將其簡單修改為解決多標簽分割問題。修改后的SegCaps 網絡僅分割左右肺部和心臟，在相同迭代次數下對特征的識別性能較低。SegCaps 網絡簡化了路由算法，減少了參數量，能分割更大圖片的同時降低了特征的識別精度。UCaps 網絡的平均分割準確率為93.21%，優于其他網絡。不同網絡的分割結果對比如圖7 所示。

表1 不同網絡的分割準確率對比Table 1 Segmentation accuracy comparison among different networks

圖7 不同網絡的分割結果對比Fig.7 Segmentation results comparison among different networks

在4 000次訓練下，不同網絡的測試結果對比如圖8所示。在訓練初期，幾種網絡模型都快速收斂，但在將每個分割部位進行準確度測算時，由于左右鎖骨部分與雙肺重疊，因此無法定位。雙肺分割精度收斂快，因心臟在圖像中的位置關系不準確，導致分割精度較低。與其他網絡相比，使用EM 路由的UCaps 網絡能在較短時間內通過位置關系定位鎖骨與心臟，在訓練前期具有較優的表現。當使用像素為單位整體評價準確率時，幾種網絡測試準確率差距不明顯，但對分割后的各部位單獨進行對比時，UCaps網絡表現十分優異。UCaps的膠囊網絡［21］結構保證各部分間的位置形態關系，此外，EM 路由與高斯聚類算法則保證UCaps 在很少的數據集中具有更優的表現。

圖8 4 000 次訓練后不同網絡的測試結果對比Fig.8 Test results comparison among different networks after 4 000 trainings

在不同訓練次數下不同網絡的右肺和右鎖骨分割準確率對比如表2 與表3 所示。從表2 和表3 可以看出，在訓練初期各網絡的右肺分割準確率較高，但在訓練初期不同網絡對右鎖骨分割時的準確率差距較大。結合EM 路由的UCaps 網絡收斂速度遠快于其他網絡模型。

表2 不同網絡的右肺分割準確率對比Table 2 Segmentation accuracy of right lung comparison among different networks %

表3 不同網絡的右鎖骨分割準確率對比Table 3 Segmentation accuracy of right clavicle comparison among different networks %

4 結束語

本文提出多標簽語義分割網絡UCaps，通過設計適用于膠囊網絡的上采樣膠囊算法，補充高層膠囊與低層膠囊的特征信息。在此基礎上，將高斯混合模型與EM 路由算法相結合，聚合底層特征對高層特征的推導過程。實驗結果表明，相比U-Net、SegCaps、MaVec-Caps 網絡，UCaps 網絡具有較快的收斂速度和較高的分割精度，在邊緣輪廓上的分割結果更清晰。后續將在圖像分割領域中通過使用金字塔結構、引入注意力機制等方法，優化膠囊網絡參數，進一步提高膠囊網絡解決不變性與同變性問題的能力。