基于卷積神經網絡的內窺鏡圖像分類

項詩雨,魏利勝

(安徽工程大學 電氣工程學院,安徽 蕪湖 241000)

近年來,我國結直腸癌發病率與死亡率都呈現增長趨勢。據2018年中國癌癥統計報告顯示,我國結直腸癌新發病例37.6萬例,死亡病例19.1萬例,結直腸癌嚴重危害并影響著我國公民健康[1]。潰瘍性結腸炎與息肉被認為是導致結腸癌的高危因素之一[2]。無線膠囊內窺鏡是檢查腸道的常用工具,每次檢查會產生數以萬計的圖像。腸道疾病的診斷需要醫生高強度地分辨大量視頻幀,即使是經驗豐富的醫生也會因為疲勞導致漏檢誤檢,因此借助計算機診斷系統來降低誤診率是非常有必要的[3]。

隨著深度學習的發展,很多基于深度學習的腸胃病變分類方法被提出來。阿依木克地斯等[4]將腸鏡圖像輸入卷積神經網絡(CNN),采用端到端的訓練方式,最高達到95.27%的息肉分類準確率;Zeng等[5]將Xception、ResNet和DenseNet三個預訓練模型融合對潰瘍性直腸炎圖像進行分類,實現了97.93%的分類準確率;Alaskar等[6]使用GoogLeNet和AlexNet模型在ImageNet數據集上進行預訓練,使用網絡參數的最佳組合高精度地分類出潰瘍。這些方法可以有效地分類內窺鏡圖像。然而,為了更高效地提取圖像特征,研究人員試圖將注意力機制與卷積神經網絡結合起來。注意力機制模型被廣泛應用于各種深度學習任務,如分割、分類等領域。它可以快速掃描圖像,類似于人眼觀察物體的方式,捕獲需要鎖定的區域。目前已經開發出具有不同結構和特征的注意力機制,將其應用于圖像處理任務中并取得了優異的結果[7-9]。鞏稼民等[10]對殘差網絡ResNet進行改進并嵌入注意力機制SE模塊實現慢性萎縮性胃炎的分類;Cao等[11]將通道注意機制集成到殘差網絡ResNet中以提取圖像特征,然后采用基于改進的SVM(支持向量機)分類器進行息肉圖像的分類,達到了98.4%的分類準確率。

本文在以上研究的基礎上,結合數據增強和遷移學習策略提出一種基于ConvNeXt網絡模型的分類方法,并在該模型中引入注意力機制CA(CoordAttention),同時使用Polyloss損失函數優化模型,實現小樣本下的潰瘍和息肉內窺鏡圖像的精確分類。

1 改進ConvNeXt模型的內窺鏡圖像分類

為了提高內窺鏡圖像分類準確率,采用了一種微調ConvNeXt模型的算法,具體流程如圖1所示。由圖1可知,改進ConvNeXt模型圖像分類主要工作為:首先需對內窺鏡圖像進行預處理操作,包括調整數據集圖像分辨率大小和修復圖像的光斑;其次內窺鏡圖像進行數據擴增,包括內窺鏡圖像水平垂直翻轉;然后結合遷移學習在引入注意力機制CA的ConvNeXt模型上進行訓練,并在驗證集上對模型的性能進行評估;最后,通過預測腳本對輸入的圖像進行分類。

圖1 基于改進ConvNeXt模型的圖像分類流程圖

1.1 圖像預處理

首先,裁剪數據集圖像。由于內窺鏡圖像的周圍存在黑框,這一部分沒有包含圖像的特性,并且會影響光斑修復操作的結果,所以將其統一裁剪掉。

其次,修復數據集圖像光斑。針對內窺鏡采集照片時會受到光照的影響產生光斑,從而影響網絡的注意力導致誤檢的問題,提出一種光斑修復算法,先對內窺鏡圖像的光斑部分進行檢測并生成mask掩膜,然后去除光斑并用周圍的像素值代替。具體算法步驟:①圖像閾值分割生成mask掩膜,本文使用的閾值范圍是200～255;②使用opencv自帶的圖像修復函數INPAINT_TELEA實現圖像的修復,其中INPAINT_TELEA函數會對位于點附近、邊界法線附近和邊界輪廓上的像素賦予更多權重,一個像素完成修復以后,它將使用快速行進的方法移動到下一個最近的像素進行修復。

最后,調整數據集圖像的分辨率。由于數據集中圖片的分辨率不一致,對模型的訓練有一定的影響。因此,為了保持網絡輸入的一致性,將內窺鏡圖像的分辨率統一調整為224×224像素。通過上述圖像預處理方法后的圖像與原始圖像之間的比較如圖2所示。從圖2圖像預處理前后的對比可以看出,使用所提出的預處理方法對內窺鏡圖像進行預處理后,圖像的光斑被消除,圖像的病變更加明顯,提高了后續網絡模型學習數據集的效率,并在一定程度上提高了算法的魯棒性。

圖2 圖像預處理

1.2 圖像增強

在深度學習中,相同訓練條件下,大的樣本量更有利于模型的訓練。由于本文所使用數據集比較小,為了防止過擬合,使用了一種數據增強方法。該方法隨機采用水平翻轉、垂直翻轉以及水平垂直翻轉相結合這3種數據增強策略,保證擴增圖像真實性的同時提高模型的泛化性能和魯棒性。數據增強效果圖如圖3所示。

圖3 數據增強

1.3 改進的ConvNeXt模型

(1)預訓練模型。采用ConvNeXt純卷積神經網絡作為預訓練模型,ConvNeXt從ResNet出發,依次從宏觀設計、深度可分離卷積、逆瓶頸層、大卷積核、細節設計這5個角度依次借鑒Swin Transformer思想,然后在ImageNet上進行訓練和評估,最終得到ConvNeXt的核心結構[12]。ConvNeXt有5個版本:ConvNeXt-T、ConvNeXt-S、ConvNeXt-B、ConvNeXt-L、ConvNeXtXL。本文將選擇ConvNeXt-T作為預訓練模型,ConvNeXt-T結構圖如圖4所示。

圖4 ConvNeXt-T模型結構圖

(2)改進的模型。雖然CNN可以獲得內窺鏡圖像的特征信息,但無法提取一些關鍵信息。傳統的注意力機制有SENet、CBAM等,其中SENet主要關注通道上的信息而忽略了位置信息;CBAM將通道注意力機制和空間注意力機制進行結合,現有實驗表明,其性能優于SENet[13],但是CBAM對病變的關注不如CA注意力機制,CA不僅捕獲跨通道的信息,還能捕獲方向感知和位置敏感信息,這有助于模型更準確地定位和識別感興趣的對象,從而實現更高的分類準確率。因此,本文在ConvNeXt網絡結構中添加了CA[14]注意力機制。注意力機制模塊作為一個即插即用的模塊,可以集成在任何特征圖后面,但是由于CA模塊的權值具有隨機初始化的特點,如果加在網絡的主干部分會破壞網絡主干部分的權值,并且網絡所使用的預訓練權重會失去作用,所以改進的模型將CA模塊添加到ConvNeXt的非主干部分。將ConvNeXt網絡的最后一個ConvNeXt Block作為輸入特征輸入CA模塊,CA模塊得到的輸出結果輸入到ConvNeXt網絡的全局池化層進行后續的分類。CA集成到ConvNeXt網絡中的結構如圖5所示。其中,C代表通道數;H代表高度;W代表寬度;r是用于控制塊大小的縮減率。圖5中CA模塊首先對輸入特征沿著水平和垂直方向進行1維平均全局池化,然后拼接起來進行卷積;其次,經過BN+非線性激活函數后,將特征圖分割開來分別進行卷積,同時關注水平和垂直方向,然后進入Sigmoid函數;最后得到的兩個注意力圖就能夠很好地反映出感興趣的對象是否存在于相應的行和列中,能夠準確地定位出目標對象的位置。

圖5 CA集成到ConvNeXt網絡中的結構圖

(3)遷移學習。由于醫學圖像數據集較小,訓練深度學習模型比較困難并且訓練后模型的泛化能力較弱,因此使用遷移學習的方法。遷移學習首先是在存在大量數據的源域(通常是ImageNet數據集)上進行訓練得到預訓練權重,然后通過在相關但不同的目標域上訓練并且微調預訓練權重來實現的。本文使用ConvNeXt-T網絡在ImageNet-1k數據集上的預訓練權重。首先,由于添加了注意力機制,網絡最后3層(Global Avg Pooling,Layer Norm,Linear)的預訓練權重失去作用,因此要將其刪除。其次,由于ImageNet-1k數據集有1 000個類別,本文使用的內窺鏡數據集只有3個類別,所以需要刪除最后一個全連接層,取而代之的是適合內窺鏡圖像分類的全連接層。最后,訓練網絡所有的權重。微調后網絡最后幾層如圖6所示。

圖6 微調后的網絡最后幾層結構圖

(4)損失函數。為了進一步優化模型,將Polyloss損失函數應用在所提出的模型中。Polyloss是一種將分類損失函數加入泰勒展開式的損失函數。Polyloss損失函數由標準交叉熵損失函數修改而來。交叉熵損失函數如式(1)所示:

lossce=-ylogy′-(1-y)log(1-y′),

(1)

式中,y是標簽,y′是預測值。

交叉熵損失函數的泰勒展開式如式(2)所示。

(2)

式中,Pt表示目標標簽預測的概率。

Polyloss損失函數僅僅修改了交叉熵損失中的第一個多項式系數,如式(3)所示:

(3)

式中,ε1是交叉熵損失中的第一個多項式系數,為了達到最佳效果,需要針對不同的任務和數據調整這個值,最佳值可以通過超參數調整找到。所提模型通過實驗選擇最優ε1=1。

1.4 評價指標

本文選用準確度、精度、召回率和F1分數4個指標來評估和分析內窺鏡圖像分類性能,其中,準確性(Accuracy)表示模型預測所有樣本中被正確分類樣本所占比例;精度(Precision)是指模型預測的所有Positive中預測正確的比例;召回率(Recall)表示所有真實Positive中預測正確的比例;F1分數表示精確率和召回率的調和平均數。各指標計算式如式(4)～(7)所示:

(4)

(5)

(6)

(7)

式中,TP表示被模型預測為正類的正樣本;TN表示被模型預測為負類的負樣本;FP表示被模型預測為正類的負樣本;FN表示被模型預測為負類的正樣本。

2 實驗驗證

為驗證基于改進ConvNeXt模型圖像分類算法的可行性和有效性,使用Pytorch深度學習框架,采用的硬件配置為Windows 11操作系統,AMD R7 CPU,GTX3060顯卡,6G顯存。首先利用Hyper Kvasir數據集對所提方法進行了訓練和驗證,并利用收斂速度較快的AdamW優化器對模型進行了優化。模型訓練迭代次數設置為50,批量大小設置為4,AdamW優化器的初始學習速率設置為0.000 5,weight_decay設置為0.005。

本文使用的數據集是公開的Hyper Kvasir數據集中的部分圖像,包含1 500張圖像,其中正常腸道、息肉、潰瘍性結腸炎各500張。此數據集較小,這對于模型的訓練更具有挑戰性。其中20%作為驗證集,80%作為訓練集。為了防止過擬合,對訓練集使用數據增強的方法擴增為正常腸道、息肉、潰瘍性結腸炎圖像各600張,共1 800張,用來訓練模型。

使用Cross Entropy loss、Focal loss以及Polyloss損失函數分別訓練所提出的模型,損失函數曲線如圖7所示。從圖7可以看出,Polyloss損失函數大概在迭代41次左右趨于穩定,Cross Entropy loss損失函數以及Focal loss損失函數在迭代45次左右趨于穩定。并且,相對于Cross Entropyloss以及Focal loss損失函數,Polyloss損失函數在訓練時損失變化更加穩定,魯棒性更強。

圖7 不同損失函數的損失曲線

為了進一步驗證改進后模型的優越性,在相同訓練條件下,使用ConvNeXt原模型、添加注意力機制CBAM的模型,以及添加CA注意力機制的模型進行對比實驗,其結果如表1所示。由表1可以看出,添加CA注意力機制的模型的準確率比原模型準確率高1.6%,比添加注意力機制CBAM的準確率高1.0%,并且使用本文模型在驗證集上能夠準確識別正常腸道,沒有把病變檢測為正常的現象,不會導致漏檢,具有一定的臨床應用潛力。3個模型的Grad-CAM[15]可視化如圖8所示。由圖8可以看出原始ConvNeXt模型以及添加CBAM注意力機制模型的特征圖略顯粗糙,專注于更多不相關的區域。相比之下,改進后的模型可以更準確地聚焦病變在內窺鏡圖像中的位置,在遇到大面積的潰瘍性結腸炎時,也能相對準確的關注所有病變區域。

表1 添加不同注意力機制后的比較

圖8 Grad-CAM可視化

為了測試該模型的分類性能,將經典網絡模型ResNet34、MobileNetV3、EfficientNetV2與改進后的模型進行了實驗分析。在相同條件下進行訓練,實驗結果如表2所示。由表2可以看出,改進后的模型相對于ResNet34、MobileNetV3、EfficientNetV2準確率分別提升1.6%、5.6%、4.5%,從每種病變的Precision、Recall以及F1-Score指標來看,改進的模型基本優于其他模型,說明改進后模型在內窺鏡圖像分類上的優越性。

表2 與經典深度學習模型對比試驗

為了進一步證明該模型的普遍適用性,將該方法應用于另一個內窺鏡圖像分類數據集Kvasir。該數據集有8個類別,每個類別1 000張圖片共8 000張,其中80%作為訓練集20%作為驗證集。使用各類評估指標的宏平均作為總數據集的評價指標。并且與其他文獻中內窺鏡圖像分類的深度學習模型進行了比較,結果如表3所示。在表3中,使用了與文獻中其他比較算法相同的數據集,不難看出,該方法在4個常用評價指標中取得了最佳效果。其中,本方法的準確率相對于文獻[16]～[18]分別提升3.5%、0.9%、3.7%。驗證了該方法的有效性和可行性,并表明改進方法在內窺鏡圖像分類領域具有良好的前景。

表3 與其他現有方法對比試驗

3 結論

本文研究了一種結合注意力機制的卷積神經網絡的內窺鏡圖像分類方法,通過選擇ConvNeXt模型作為預訓練模型,結合遷移學習微調網絡模型結構,并添加注意力機制構成一種端到端的分類模型,同時在數據預處理階段探討一種光斑修復方法對內窺鏡圖像高光部分進行修復。實驗結果表明,改進后模型的分類性能有明顯提高,為潰瘍性結腸炎以及息肉的診斷提供了幫助。但是所提模型只對病變進行分類,沒有實現病變精準定位,如何實現內窺鏡息肉以及潰瘍位置的精準定位將是接下來的重點研究工作。