使用自適應(yīng)條形采樣和雙分支Transformer的DA-Net

文章目錄

前言

概述

網(wǎng)絡(luò)簡(jiǎn)介

DBTM：Local Patches Meet Global Context

ASUB block

實(shí)驗(yàn)

討論

前言

這是 MICCAI 2022 上的第三篇閱讀筆記了，之前兩篇也都可以在 GiantPandaCV 公眾號(hào)搜索到。如下圖所示，目前的視網(wǎng)膜血管分割方法按照輸入數(shù)據(jù)劃分有兩類(lèi)：image-level 和 patches-level，每一種方法都有自己的優(yōu)勢(shì)，如何將兩者結(jié)合起來(lái)是一個(gè)需要去解決的問(wèn)題，這也是 DA-Net 這篇文章的貢獻(xiàn)之一。此外，這篇文章還提出了一個(gè)自適應(yīng)的條狀 Upsampling Block，我們會(huì)在后面展開(kāi)介紹。

概述

目前的視網(wǎng)膜血管分割方法根據(jù)輸入類(lèi)型大致分為 image-level 和 patches-level 方法，為了從這兩種輸入形式中受益，這篇文章引入了一個(gè)雙分支 Transformer 模塊，被叫做 DBTM，它可以同時(shí)利用 patches-level 的本地信息和 image-level 的全局上下文信息。視網(wǎng)膜血管跨度長(zhǎng)、細(xì)且呈條狀分布，傳統(tǒng)的方形卷積核表現(xiàn)不佳，也是為了更好地捕獲這部分的上下文信息，進(jìn)一步設(shè)計(jì)了一個(gè)自適應(yīng)條狀 Upsampling Block，被叫做 ASUB，以適應(yīng)視網(wǎng)膜血管的條狀分布。

網(wǎng)絡(luò)簡(jiǎn)介

下圖是 DA-Net 的整體結(jié)構(gòu)。共享 encoder 包含五個(gè)卷積塊，DBTM 在 encoder 之后，最后是帶 ASUB 的 decoder。首先，原眼底圖像很常規(guī)的被分成 N^2 個(gè) patches，N 為 patch 的大小，除此之外，將原眼底圖像也下采樣 N 倍，但是不做裁剪。將它們一起送入共享 encoder，分別得到相應(yīng)的特征圖 F(i) 和 F′，這里的共享指的是兩個(gè) encoder 分支的權(quán)重共享（那么你可以把它簡(jiǎn)單理解為用同一個(gè)卷積核掃描 N^1+1 個(gè) patches，只不過(guò)其中 1 這個(gè) patch 是完整的圖像），兩個(gè)分支可以通過(guò)合并批次并行操作，這意味著輸入圖像的編碼可以在一次推理中完成，無(wú)需增加額外的參數(shù)和時(shí)間消耗。隨后，這兩個(gè)分支的輸出通過(guò) DBTM 進(jìn)行通信，DBTM 可以向每個(gè)補(bǔ)丁廣播長(zhǎng)距離的全局信息。U 型網(wǎng)絡(luò)中間還有普通的跨層連接，最后，再通過(guò) ASUB 的 decoder后，得到預(yù)測(cè)的分割結(jié)果。

請(qǐng)?zhí)砑訄D片描述

DBTM：Local Patches Meet Global Context

下面兩部分，我們分別對(duì) DBTM 和 AUSB block 展開(kāi)介紹。

首先，將經(jīng)過(guò) flatten 和投影的特征圖 F(i) 和 F′ 作為輸入 tokens ，其中加入訓(xùn)練過(guò)的 position embeddings 以保留位置信息。然后，如下圖所示，輸入 tokens 被送入 Transformer Layer。不同的是，設(shè)計(jì)了一個(gè)特殊的 self-then-cross 的 pipeline，將兩個(gè)分支的輸入混合起來(lái)，稱(chēng)為雙分支 Transformer 模塊（看網(wǎng)絡(luò)簡(jiǎn)介中的圖）。第一個(gè) Transformer Layer 作為 Q，第二個(gè) Transformer Layer 作為 K 和 V。具體來(lái)說(shuō)，首先，這兩個(gè)分支的輸入標(biāo)記分別用自注意機(jī)制模擬 image-level 和 patches-level 的長(zhǎng)距離依賴(lài)。然后，交叉注意機(jī)制被用于兩個(gè)分支的 tokens 之間的通信。在交叉注意機(jī)制中，將 patches-level 的標(biāo)記表示為查詢 Q， image-level 分支的標(biāo)記表示為下圖中多頭自我注意（MSA）層的鍵 Q 值 V。整體設(shè)計(jì)是很簡(jiǎn)單的，實(shí)現(xiàn)了”Local Patches Meet Global Context“。

ASUB block

視網(wǎng)膜血管的一些固有特征導(dǎo)致了其分割困難，比如視網(wǎng)膜血管的分支很細(xì)，邊界很難區(qū)分，而且視網(wǎng)膜血管之間的關(guān)系很復(fù)雜。在這些情況下，視網(wǎng)膜血管周?chē)谋尘靶畔?duì)視網(wǎng)膜血管的分割至關(guān)重要。如下圖所示，傳統(tǒng)的方形卷積核在正常的上采樣塊中不能很好地捕捉線性特征，并且不可避免地引入了來(lái)自鄰近的不相關(guān)信息。為了更好地收集視網(wǎng)膜血管周?chē)谋尘靶畔ⅲ岢隽?Adaptive Strip Upsampling Block（ASUB），它適合于長(zhǎng)而細(xì)的視網(wǎng)膜血管分布。

請(qǐng)?zhí)砑訄D片描述

粗略看一下，在 (c) 中，一共有四種類(lèi)型的條狀卷積核，捕捉水平（S1）、垂直（S2）、左對(duì)角線（S3）和右對(duì)角線（S4）方向上的信息。接下來(lái)，我們仔細(xì)分析下 ASUB 的思路，首先，使用一個(gè) 1×1 的 Conv 來(lái)將特征圖的維度減半，以減少計(jì)算成本。然后，利用四個(gè)帶狀卷積來(lái)捕捉來(lái)自不同方向的上下文信息。此外，做全局平均池化（GAP）來(lái)獲得通道維度的特征圖。在特征圖的通道維度上獲得特征向量，并使用全連接層來(lái)學(xué)習(xí)每個(gè)帶狀卷積的通道方向的注意立向量。之后，應(yīng)用萬(wàn)能的 softmax 來(lái)產(chǎn)生通道融合權(quán)重Wi , i∈{1, 2, 3, 4}。最后，我們用學(xué)到的自適應(yīng)權(quán)重對(duì)每個(gè)帶狀卷積 Fi 的輸出進(jìn)行加權(quán)，得到特征圖，特征圖是 4 個(gè) Fi*Wi 求和。最后用 1×1 的 Conv 恢復(fù)維度，得到最終輸出 Foutput。同時(shí)，這部分是會(huì)增加網(wǎng)絡(luò)學(xué)習(xí)負(fù)擔(dān)的。

請(qǐng)?zhí)砑訄D片描述

實(shí)驗(yàn)

首先是和其他 SOTA 方法的比較，包括 image-level 和 patches-level 兩種，如下表。

接下來(lái)是消融實(shí)驗(yàn)的部分，其中的 Baseline 指 U-Net。注意到，F(xiàn)LOPs 和 參數(shù)量的增加是可以接受的。

討論

其實(shí) ASUB 設(shè)置的條形采樣方向也不一定與一些小血管的方向完全一致，這是可以進(jìn)一步改進(jìn)的地方。比如說(shuō)嘗試可變形卷積（Deformable ConvNetsV2）的方式。