CVPR 2023 中的領(lǐng)域適應(yīng)：用于切片方向連續(xù)的無監(jiān)督跨模態(tài)醫(yī)學(xué)圖像分割

CVPR 2023 中的領(lǐng)域適應(yīng)：用于切片方向連續(xù)的無監(jiān)督跨模態(tài)醫(yī)學(xué)圖像分割

前言

我們已經(jīng)介紹過 3 篇 CVPR 中的典型領(lǐng)域適應(yīng)工作，他們?nèi)际?TTA（Test-Time Adaptation）的 settings，而這次要介紹的文章是 UDA（Unsupervised domain adaptation）的 setting。之前的三篇文章分別是：

CoTTA

EcoTTA

DIGA

在這篇文章中，提出了 SDC-UDA，一種簡(jiǎn)單而有效的用于連續(xù)切片方向的跨模態(tài)醫(yī)學(xué)圖像分割的體積型 UDA 框架，它結(jié)合了切片內(nèi)和切片間自注意力圖像轉(zhuǎn)換、不確定性約束的偽標(biāo)簽優(yōu)化和體積型自訓(xùn)練。與以前的醫(yī)學(xué)圖像分割 UDA 方法不同之處在于它可以獲得切片方向上的連續(xù)分割（這一點(diǎn)有點(diǎn)重要，因?yàn)橥R床上都是一個(gè) 3D 數(shù)據(jù)，而直接處理 3D 數(shù)據(jù)又需要很大的計(jì)算資源），從而確保更高的準(zhǔn)確性和臨床實(shí)踐中的潛力。

如上圖所示，以前的醫(yī)學(xué)圖像分割 UDA 方法大多采用 2D UDA，當(dāng)將預(yù)測(cè)堆疊在一起時(shí)，會(huì)導(dǎo)致切片方向上的預(yù)測(cè)不一致。SDC-UDA 在翻譯和分割過程中考慮了體積信息，從而改善了分割結(jié)果在切片方向上的連續(xù)性，可以看到在圖的最右側(cè)，下面方法的 Dice 值在切片方向上是穩(wěn)定的。

此外，我們?nèi)闹刑岬降摹绑w積”這個(gè)詞，可以理解為 3D 數(shù)據(jù)。

體積型 UDA 框架概述

如下圖所示，SDC-UDA 大致有五個(gè)步驟，從 stage 1 到 stage 5：

stage 1：帶有片內(nèi)和片間注意力的對(duì)抗學(xué)習(xí)過程，這一步是 stage 2 的基礎(chǔ)，stage 2 是該步驟的上半部分。后面會(huì)單獨(dú)用一個(gè)小節(jié)介紹。

stage 2：target 模態(tài)數(shù)據(jù)生成，假如 source 數(shù)據(jù)模態(tài)是 MRI，那么在這個(gè)步驟我們會(huì)得到 3D 的 CT 和對(duì)應(yīng)的 label。

stage 3：把生成的 target 數(shù)據(jù)和 label 送入到教師網(wǎng)絡(luò)訓(xùn)練。

stage 4：將真實(shí)的不帶標(biāo)簽的 target 數(shù)據(jù)輸入到 stage 3 的教師網(wǎng)絡(luò)得到偽標(biāo)簽，并通過不確定性抑制優(yōu)化偽標(biāo)簽。

stage 5：將生成的 target 數(shù)據(jù)、真實(shí) target 數(shù)據(jù)和他們的標(biāo)簽用于優(yōu)化學(xué)生網(wǎng)絡(luò)，最終的預(yù)測(cè)也是在學(xué)生網(wǎng)絡(luò)上。

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

具體實(shí)現(xiàn)

Unpaired 圖像轉(zhuǎn)換

先前的 2D UDA 方法將 3D 體積分割成 2D 切片，并在之后將它們的轉(zhuǎn)換重新堆疊成 3D 體積。由于切片是單獨(dú)處理的，重新構(gòu)建轉(zhuǎn)換后的體積通常需要額外的后處理，如切片方向插值，但這仍然無法完全解決切片方向不連續(xù)等問題。為了解決 2D 方法缺乏對(duì)體積性質(zhì)的考慮和 3D 方法的優(yōu)化效率問題，這篇文章提出了一種簡(jiǎn)單而有效的像素級(jí)領(lǐng)域轉(zhuǎn)換方法，用于醫(yī)學(xué)圖像體積數(shù)據(jù)，通過使用切片內(nèi)部和切片間自注意力模塊將一組源域圖像轉(zhuǎn)換為目標(biāo)域圖像。與先前的 2D 方法只在單個(gè)切片內(nèi)進(jìn)行轉(zhuǎn)換，而這篇文章的方法利用了切片方向上相鄰切片的信息。這類似于最近在視頻處理中的進(jìn)展，它利用了幀內(nèi)部和幀之間的信息。與需要昂貴計(jì)算成本的 3D 方法相比，不需要大量計(jì)算（下采樣）。

如上圖所示，首先我們將一個(gè) 3D MRI 數(shù)據(jù)裁剪出 3 張切片，輸入到 CNN 的 encoder中，encoder 的輸出是三張切片的 feature maps，即在通道維度上被卷積。然后我們?cè)陂L(zhǎng)和寬的方向上裁剪 patches，這樣會(huì)得到若干個(gè) patch 塊，輸入到帶有片內(nèi)和片間的切片注意力模塊中。這個(gè)注意力模塊就是很普通的多頭注意力、殘差和 FFN 的兩次組合。最后我們做相反過程的 decoder，這時(shí)生成的圖像應(yīng)該是 target 模態(tài)的。為了方便理解，可以再去看看我們?cè)谏弦还?jié)提到的 stage 1，對(duì)應(yīng) stage 1 的上半部分。

stage 1 除了包括上面提到的這個(gè)過程，還包括重建的反過程（下半部分），這樣我們才能計(jì)算一致性的 loss，同時(shí)利用對(duì)抗學(xué)習(xí)的判別器，完成自監(jiān)督的訓(xùn)練。

體積自訓(xùn)練和偽標(biāo)簽優(yōu)化

我們已經(jīng)介紹了概述中第一個(gè) stage，這一節(jié)對(duì)應(yīng)后面三個(gè) stage。

通過從源域轉(zhuǎn)換的合成數(shù)據(jù) x?t 和注釋 ys（即帶標(biāo)簽的合成數(shù)據(jù)集），我們首先訓(xùn)練一個(gè)教師分割網(wǎng)絡(luò) teacher，該網(wǎng)絡(luò)最小化分割損失：

訓(xùn)練完教師模型，可以通過將真實(shí)的目標(biāo)域數(shù)據(jù) xt 傳遞給訓(xùn)練好的分割模型 teacher，獲取未標(biāo)記真實(shí)數(shù)據(jù)的偽標(biāo)簽 y?t。

由于 teacher 預(yù)測(cè)出的偽標(biāo)簽是噪聲標(biāo)簽，必須對(duì)其進(jìn)行改進(jìn)，以提高準(zhǔn)確性并引導(dǎo)自訓(xùn)練朝更好的方向發(fā)展。這篇文章設(shè)計(jì)了一種增強(qiáng)敏感性（SE）和特異性（SP）的偽標(biāo)簽改進(jìn)模塊，該模塊基于圖像強(qiáng)度、當(dāng)前偽標(biāo)簽和不確定性區(qū)域（高于閾值）來改進(jìn)偽標(biāo)簽。

通過預(yù)測(cè)出的偽標(biāo)簽，計(jì)算與每個(gè)類別相對(duì)應(yīng)的不確定性（即熵）圖：

其中 p 是每個(gè)類別的輸出概率圖。為了增強(qiáng)偽標(biāo)簽的敏感性，檢測(cè)超出偽標(biāo)簽范圍的高度不確定的區(qū)域。然后，如果該區(qū)域中的像素強(qiáng)度在當(dāng)前偽標(biāo)簽包含的圖像強(qiáng)度的某個(gè)范圍內(nèi)，該區(qū)域?qū)⒈话閭螛?biāo)簽的一部分。該公式可以表示為：

其中 分別表示目標(biāo)域圖像、偽標(biāo)簽、改進(jìn)的偽標(biāo)簽和裁剪了高不確定性區(qū)域掩碼。該方法基于假設(shè)：在醫(yī)學(xué)圖像中，具有相似強(qiáng)度且相互接近的像素很可能屬于同一類別。

為了增強(qiáng)偽標(biāo)簽的特異性，也是檢測(cè)偽標(biāo)簽范圍內(nèi)的高度不確定的區(qū)域。區(qū)別是，如果該區(qū)域中的像素強(qiáng)度不在當(dāng)前偽標(biāo)簽包含的圖像強(qiáng)度的某個(gè)范圍內(nèi)，則將其從當(dāng)前偽標(biāo)簽中排除?？梢员硎緸椋?/p>

上面這個(gè)流程，文章中給出了圖示如下，有助于理解這個(gè)流程：

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

在 stage 5 中，合成的 target scans 與真實(shí) target scans 存在分布差異。這篇文章將這兩種配對(duì)數(shù)據(jù)結(jié)合到自訓(xùn)練中，以最大程度地提高泛化能力，并最小化由于分布差異而引起的性能下降。把帶標(biāo)簽的合成 target scans 和帶偽標(biāo)簽的 target scans 的數(shù)據(jù)合并，訓(xùn)練一個(gè)學(xué)生分割模型 student，以最小化以下?lián)p失：

實(shí)驗(yàn)

下表是 SDC-UDA 與以前的非醫(yī)學(xué)圖像和醫(yī)學(xué)圖像 UDA 方法之間的定量結(jié)果的比較。該表包括非醫(yī)學(xué)圖像 UDA 方法（例如 cycleGan、cycada、ADVENT 和 FDA）的結(jié)果，以及最近的醫(yī)學(xué)圖像 UDA 方法（例如 SIFA 和 PSIGAN）的結(jié)果。對(duì)比發(fā)表在 TMI 2020 上的 PSIGAN 方法，DICE 指標(biāo)上提升了很多，特別是從 T1 到 T2 的跨模態(tài)設(shè)置。MRI 到 CT 也有顯著的提升。

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

可視化結(jié)果比較如下圖：

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

總結(jié)

這篇文章提出了 SDC-UDA，一種用于切片方向連續(xù)的跨模態(tài)醫(yī)學(xué)圖像分割的新型 UDA 框架。SDC-UDA 通過切片內(nèi)部和切片間的自注意力有效地轉(zhuǎn)換醫(yī)學(xué)體積，并通過利用不確定性圖，設(shè)計(jì)簡(jiǎn)單而有效的偽標(biāo)簽細(xì)化策略。通過體積級(jí)自訓(xùn)練更好地適應(yīng)目標(biāo)域。

現(xiàn)在的 SDC-UDA 框架中，只有 stage 1 是不需要訓(xùn)練 3D 圖像的，后面的過程仍然是 3D 的訓(xùn)練（可能出于準(zhǔn)確率的角度），也需要消耗更多的計(jì)算資源，其實(shí)也是可以優(yōu)化成一組堆疊切片的。

參考

https://arxiv.org/pdf/2305.11012.pdf