做醫(yī)學(xué)圖像分割的同學(xué)應(yīng)該都有同感病灶數(shù)據(jù)永遠(yuǎn)是“不夠用”的。一份高質(zhì)量的腫瘤掩膜需要影像科醫(yī)生在 CT、MRI 切片上一層一層手動(dòng)勾畫成本高、周期長(zhǎng)而且部分病灶形態(tài)本身就十分罕見——比如小尺寸轉(zhuǎn)移瘤、貼近血管生長(zhǎng)的結(jié)節(jié)、跨解剖區(qū)域的不規(guī)則病變。如果訓(xùn)練集里從未出現(xiàn)過(guò)這樣的形狀和位置模型到了真實(shí)場(chǎng)景中很容易漏檢、誤檢甚至把正常組織判成病灶。常規(guī)的緩解手段是數(shù)據(jù)增強(qiáng)旋轉(zhuǎn)、翻轉(zhuǎn)、縮放、彈性形變都是常見操作。但這些方法本質(zhì)是在“已有病灶”的基礎(chǔ)上做幾何擾動(dòng)并不具備生成“新形狀、新位置病灶”的能力。于是問(wèn)題變得很有意思能不能像捏橡皮泥一樣把兩個(gè)真實(shí)病灶的形狀、位置和紋理做一次可解釋的融合從而合成大量“介于兩者之間”的逼真病灶樣本OTLesMix 正是沿著這條思路提出的方法。它的核心不是復(fù)雜神經(jīng)網(wǎng)絡(luò)而是一套經(jīng)典數(shù)學(xué)工具最優(yōu)傳輸Optimal Transport, OT與 Wasserstein 重心Wasserstein Barycenter。這篇文章我會(huì)從問(wèn)題背景、數(shù)學(xué)原理、方法拆解、代碼復(fù)現(xiàn)思路、實(shí)驗(yàn)驗(yàn)證方式和工程落地建議幾個(gè)維度把 OTLesMix 完整地講清楚。無(wú)論你是醫(yī)學(xué)圖像方向的算法工程師還是對(duì)最優(yōu)傳輸在深度學(xué)習(xí)中的應(yīng)用感興趣的研究者這篇文章都會(huì)對(duì)你有幫助。1. 醫(yī)學(xué)圖像分割的病灶數(shù)據(jù)難題1.1 病灶數(shù)據(jù)為什么稀缺病灶分割數(shù)據(jù)的稀缺并不是單純“數(shù)據(jù)量少”而是好幾個(gè)因素疊加在一起的結(jié)果。首先是標(biāo)注成本。CTA、MRI 或者病理切片中的病灶邊界往往并不清晰不同醫(yī)生對(duì)同一病灶的勾畫也可能存在差異。要獲得一份高質(zhì)量標(biāo)注通常需要多位醫(yī)生交叉驗(yàn)證甚至病理結(jié)果對(duì)照這在時(shí)間、人力和資金上都是很大開銷。其次是隱私約束。醫(yī)學(xué)影像屬于敏感個(gè)人數(shù)據(jù)脫敏、加密、倫理審批、數(shù)據(jù)使用協(xié)議等環(huán)節(jié)都會(huì)拖慢數(shù)據(jù)收集節(jié)奏。即便拿到數(shù)據(jù)不同醫(yī)院之間的數(shù)據(jù)格式、掃描協(xié)議、設(shè)備型號(hào)差異也會(huì)造成分布偏移。最核心的問(wèn)題是長(zhǎng)尾分布。病灶在真實(shí)世界中不是均勻出現(xiàn)的比如早期小病灶、罕見位置的病灶在數(shù)據(jù)集中占比很低而數(shù)據(jù)增強(qiáng)恰恰需要在這些樣本上下更多功夫。如果只是簡(jiǎn)單復(fù)制粘貼這類樣本模型很快就會(huì)過(guò)擬合換一個(gè)姿勢(shì)、一個(gè)角度、一個(gè)位置就認(rèn)不出來(lái)。1.2 傳統(tǒng)數(shù)據(jù)增強(qiáng)為什么不夠先列一下醫(yī)學(xué)分割里用得最多的幾何增強(qiáng)增強(qiáng)方式作用局限旋轉(zhuǎn)、翻轉(zhuǎn)改變病灶方向不改變形狀本質(zhì)小病灶仍難泛化縮放改變病灶尺寸實(shí)際病變大小與解剖結(jié)構(gòu)相關(guān)不能隨意縮放彈性形變輕微扭曲形狀幅度太大易失真幅度太小不足以增加多樣性亮度對(duì)比度擾動(dòng)模擬不同掃描參數(shù)不能生成新的結(jié)構(gòu)信息Mixup / CutMix像素級(jí)或塊級(jí)混合通常不做空間對(duì)應(yīng)病灶位置錯(cuò)開時(shí)會(huì)產(chǎn)生重影Mixup、CutMix 這類通用增強(qiáng)方法在自然圖像分類上效果很好但直接搬到醫(yī)學(xué)圖像分割時(shí)有一個(gè)明顯問(wèn)題病灶不是“一塊普通貼紙”它有自身的形狀組織、紋理統(tǒng)計(jì)和空間語(yǔ)義。當(dāng)兩個(gè)病灶位置距離較遠(yuǎn)時(shí)線性插值會(huì)產(chǎn)生模糊和重影當(dāng)你把病灶 A 直接復(fù)制到病灶 B 的位置時(shí)兩個(gè)病灶邊緣之間的過(guò)渡又非常突兀模型學(xué)到的更多是“貼圖偽影”而不是有意義的病理結(jié)構(gòu)。1.3 合成病灶的兩條技術(shù)路線為了讓訓(xùn)練集覆蓋更多形狀和位置研究者通常走兩條路。一條是生成模型路線典型代表是 GAN 和擴(kuò)散模型。GAN 可以生成非常逼真的合成病灶但訓(xùn)練不穩(wěn)定、需要大量真實(shí)樣本作為監(jiān)督、生成結(jié)果的病理語(yǔ)義難以保證。擴(kuò)散模型效果更細(xì)膩但采樣速度慢、訓(xùn)練成本高并且依然存在“生成的內(nèi)容是否符合真實(shí)病理特征”這一可解釋性問(wèn)題。另一條路線是幾何變換與圖像混合路線OTLesMix 屬于這一類。它在兩個(gè)真實(shí)病灶之間建立一個(gè)“質(zhì)量搬運(yùn)”方案讓形狀、位置和紋理可以解耦處理。這樣做的好處是合成病灶的紋理完全來(lái)自真實(shí)病灶不存在紋理漂移問(wèn)題形狀和位置則來(lái)自數(shù)學(xué)上可解釋的插值過(guò)程穩(wěn)定且可控。2. 先搞清楚最優(yōu)傳輸與 Wasserstein 重心2.1 最優(yōu)傳輸問(wèn)題從搬運(yùn)沙土說(shuō)起最優(yōu)傳輸是一個(gè)有 200 多年歷史的數(shù)學(xué)問(wèn)題最早可以追溯到法國(guó)工程師 Monge 在 1781 年提出的“搬土問(wèn)題”。想象地上有一堆沙子分布形狀為 A現(xiàn)在要求你把它搬成分布形狀為 B并且每一粒沙子從起點(diǎn)到終點(diǎn)都會(huì)產(chǎn)生搬運(yùn)成本。如果搬運(yùn)成本用兩點(diǎn)之間的距離來(lái)度量那么最優(yōu)傳輸問(wèn)題要回答的就是怎樣規(guī)劃每一粒沙子的搬運(yùn)路徑讓總成本最低這個(gè)“搬運(yùn)方案”在數(shù)學(xué)上稱為傳輸計(jì)劃transport plan記作 γ(x, y)表示有多少質(zhì)量從源位置 x 流向目標(biāo)位置 y。Kantorovich 后來(lái)對(duì)問(wèn)題做了松弛化處理允許質(zhì)量被“拆分”于是最優(yōu)傳輸問(wèn)題變成了一個(gè)線性規(guī)劃問(wèn)題min ∫ c(x, y) dγ(x, y)其中 c(x, y) 是搬運(yùn)代價(jià)函數(shù)通常取距離的平方 ||x - y||2。這個(gè)公式看起來(lái)抽象但含義非常樸素在所有可能的搬運(yùn)方案里找一個(gè)總成本最小的。在圖像處理里一張圖像的灰度值可以歸一化成一個(gè)概率分布病灶的掩膜也可以看成一個(gè)概率分布。于是最優(yōu)傳輸就提供了一種“把一個(gè)病灶變成另一個(gè)病灶”的最優(yōu)幾何對(duì)應(yīng)關(guān)系。2.2 Wasserstein 距離分布之間的距離當(dāng)最優(yōu)傳輸?shù)淖顑?yōu)總代價(jià)被定義出來(lái)后它本身也可以用來(lái)度量?jī)蓚€(gè)分布之間的距離。這就是 Wasserstein 距離。Wasserstein-p 距離的定義是W_p(μ, ν) (min ∫ ||x - y||^p dγ(x, y))^(1/p)當(dāng) p 1 時(shí)它也叫推土機(jī)距離Earth Movers Distance, EMD在圖像檢索、生成模型評(píng)估等領(lǐng)域有廣泛應(yīng)用。Wasserstein 距離有一個(gè)非常重要的幾何直覺它考慮的是“質(zhì)量分布的形狀差異”。兩個(gè)分布即使支撐集完全不重疊Wasserstein 距離也依然有明確的有限值因?yàn)橘|(zhì)量可以沿著空間距離流動(dòng)。對(duì)比 KL 散度和 JS 散度它們?cè)趦蓚€(gè)分布支撐集不重疊時(shí)可能退化為無(wú)窮大或常數(shù)導(dǎo)致梯度消失而 Wasserstein 距離不會(huì)。在 OTLesMix 中Wasserstein 距離可以用來(lái)度量?jī)蓚€(gè)病灶掩膜之間的“形狀距離”也可以作為插值過(guò)程中的一致性指標(biāo)。2.3 Wasserstein Barycenter分布的“平均”給定多個(gè)概率分布 μ?, μ?, ..., μ_N它們的 Wasserstein 重心是一個(gè)新分布 μ*使得所有分布與 μ* 之間 Wasserstein 距離的加權(quán)和最小μ* argmin_μ Σ λ_k · W_p(μ, μ_k)^p如果 λ_k 是權(quán)重且所有權(quán)重之和為 1那么 μ* 就可以理解為這些分布在“最優(yōu)傳輸意義下的平均”。這種平均和歐氏平均很不一樣。歐氏平均只是逐點(diǎn)取像素值的加權(quán)平均容易產(chǎn)生模糊、重影而 Wasserstein 重心會(huì)盡量保持分布本身的結(jié)構(gòu)——在病灶掩膜的場(chǎng)景里它傾向于生成一個(gè)“形狀介于兩個(gè)病灶之間”的掩膜而不是簡(jiǎn)單地灰度混合。舉個(gè)例子一個(gè)病灶是圓形另一個(gè)是細(xì)長(zhǎng)形它們的歐氏平均可能是一團(tuán)模糊的灰塊但 Wasserstein 重心可以生成一個(gè)從圓形到細(xì)長(zhǎng)形漸變過(guò)程中的“中間形態(tài)”這個(gè)形態(tài)仍然具有清晰的輪廓和可解釋的形狀結(jié)構(gòu)。3. OTLesMix 方法拆解形狀、位置與紋理如何被“搬運(yùn)”從命名看OTLesMix 可以拆解為 OT最優(yōu)傳輸 Les病變 Lesion Mix混合。它要解決的核心問(wèn)題是給定兩個(gè)真實(shí)病灶樣本如何合成一系列形狀、位置都不同但紋理統(tǒng)計(jì)真實(shí)可信的新病灶。3.1 病灶掩膜作為概率分布OTLesMix 首先把病灶掩膜mask看作一個(gè)二維概率分布。假設(shè)掩膜中病灶區(qū)域像素值為 1背景為 0那么將掩膜歸一化后它就是一個(gè)定義在像素平面上的離散概率分布源的病灶掩膜 μ_s質(zhì)量集中在病灶區(qū)域目標(biāo)的病灶掩膜 μ_t質(zhì)量集中在另一個(gè)病灶區(qū)域用概率分布的語(yǔ)言來(lái)描述病灶是 OTLesMix 最關(guān)鍵的一步。因?yàn)檫@樣一來(lái)形狀差異就變成了“兩個(gè)概率分布之間的差異”而形狀插值就變成了“分布之間的重心計(jì)算”。3.2 最優(yōu)傳輸映射建立空間對(duì)應(yīng)關(guān)系有了兩個(gè)概率分布下一步是計(jì)算它們之間的最優(yōu)傳輸映射。這里要說(shuō)一下“最優(yōu)傳輸映射”與“最優(yōu)傳輸計(jì)劃”的區(qū)別。最優(yōu)傳輸計(jì)劃 γ(x, y) 是一個(gè)聯(lián)合分布表示質(zhì)量從 x 到 y 的流量而在某些情況下這個(gè)計(jì)劃可以退化為一個(gè)確定性的函數(shù) T: x → y即每個(gè)源像素點(diǎn)都映射到唯一的目標(biāo)像素點(diǎn)這個(gè)函數(shù)就叫最優(yōu)傳輸映射。OTLesMix 使用最優(yōu)傳輸計(jì)劃的核心動(dòng)機(jī)是它比仿射變換、TPS 薄板樣條等傳統(tǒng)配準(zhǔn)方法更靈活能夠捕捉兩個(gè)病灶之間任意的質(zhì)量對(duì)應(yīng)關(guān)系同時(shí)保證這種對(duì)應(yīng)是“全局最優(yōu)”的不會(huì)因?yàn)榫植控澬亩a(chǎn)生扭曲。在工程實(shí)現(xiàn)中這一步通常通過(guò)計(jì)算兩個(gè)掩膜像素坐標(biāo)之間的代價(jià)矩陣然后求解線性規(guī)劃來(lái)完成也就是計(jì)算G EMD(μ_s, μ_t, M)其中 M 是像素坐標(biāo)之間的代價(jià)矩陣G 是傳輸計(jì)劃。G 本身可以理解為一張“從源病灶到目標(biāo)病灶的質(zhì)量流量圖”它標(biāo)識(shí)了源病灶每一個(gè)像素應(yīng)該搬運(yùn)到目標(biāo)病灶的哪個(gè)位置。3.3 Wasserstein 重心生成多樣中間形狀有了傳輸計(jì)劃之后OTLesMix 要做的不是直接跳到最極端的目標(biāo)形狀而是在源病灶和目標(biāo)病灶之間生成一系列中間形態(tài)。這一部分用到的正是 Wasserstein 重心?？紤]兩個(gè)掩膜 μ_s 和 μ_t取權(quán)重 λ ∈ [0, 1]計(jì)算它們的 Wasserstein 重心μ_λ argmin_μ [ (1-λ) · W_p(μ, μ_s)^p λ · W_p(μ, μ_t)^p ]當(dāng) λ 0 時(shí)重心就是源掩膜當(dāng) λ 1 時(shí)重心就是目標(biāo)掩膜當(dāng) λ 0.5 時(shí)重心是“中間形狀”。由于 Wasserstein 重心本身是概率分布因此合成掩膜依然保留著清晰的形態(tài)結(jié)構(gòu)不會(huì)像線性插值那樣出現(xiàn)模糊和重影。通過(guò)改變 λ 的取值可以生成一系列形狀平滑變化的病灶掩膜這就是“Diverse Shapes”的來(lái)源。3.4 紋理保持與位置遷移單純生成形狀還不夠合成病灶必須帶上真實(shí)的紋理灰度統(tǒng)計(jì)特征才可用于模型訓(xùn)練。OTLesMix 的做法是用最優(yōu)傳輸計(jì)劃把源病灶的圖像紋理“搬運(yùn)”到目標(biāo)病灶的位置上去。具體來(lái)說(shuō)給定傳輸計(jì)劃 G對(duì)于目標(biāo)位置的每一個(gè)像素我們根據(jù) G 中對(duì)應(yīng)的權(quán)重從源病灶圖像里加權(quán)采樣像素值得到一張“遷移動(dòng)態(tài)紋理圖”。這樣整個(gè)方法就把病灶的“形狀”和“紋理”解耦了形狀來(lái)自 Wasserstein 重心插值紋理來(lái)自真實(shí)源病灶的傳輸搬運(yùn)位置來(lái)自最優(yōu)傳輸映射所建立的空間對(duì)應(yīng)關(guān)系通過(guò)改變?cè)床≡睢⒛繕?biāo)病灶的組合以及權(quán)重 λ可以生成“形狀像 A、紋理像 B、位置接近 A/B 之間”的大量合成病灶。病灶位置的多樣性主要來(lái)自兩個(gè)方面一是源病灶和目標(biāo)病灶本身位置不同傳輸計(jì)劃會(huì)把紋理搬運(yùn)到新的位置二是在多個(gè)真實(shí)病灶之間兩兩組合時(shí)會(huì)形成大量沒有在原始數(shù)據(jù)中出現(xiàn)過(guò)的位置分布。3.5 合成病灶如何融入訓(xùn)練數(shù)據(jù)在實(shí)際訓(xùn)練時(shí)OTLesMix 通常作為一個(gè)在線數(shù)據(jù)增強(qiáng)模塊接入分割訓(xùn)練管線。流程大致如下訓(xùn)練集中維護(hù)一個(gè)“真實(shí)病灶庫(kù)”每個(gè)樣本包含病灶小塊和對(duì)應(yīng)掩膜。每個(gè)訓(xùn)練 step 隨機(jī)抽取兩個(gè)病灶樣本。用 OTLesMix 合成新的病灶 patch。將合成 patch 粘貼到訓(xùn)練圖像的隨機(jī)位置或通過(guò) OT 映射得到的合理位置。合成 patch 對(duì)應(yīng)的掩膜直接作為訓(xùn)練標(biāo)簽。通過(guò)這種方式訓(xùn)練集每次迭代都有新的合成病灶出現(xiàn)模型能夠見到更多樣的形狀和位置組合。4. 代碼復(fù)現(xiàn)思路用 Python 實(shí)現(xiàn)核心模塊下面給出一個(gè)基于 POTPython Optimal Transport庫(kù)的示例實(shí)現(xiàn)。需要提前說(shuō)明這部分代碼是為了幫助你理解 OTLesMix 的核心思路并不是論文官方實(shí)現(xiàn)的完整代碼。實(shí)際復(fù)現(xiàn)時(shí)需要根據(jù)你自己的數(shù)據(jù)格式、病灶尺寸和訓(xùn)練框架進(jìn)行調(diào)整。4.1 環(huán)境準(zhǔn)備與依賴建議使用以下環(huán)境Python 3.8 PyTorch 1.10訓(xùn)練分割模型用 POT 0.9.0最優(yōu)傳輸計(jì)算 NumPy OpenCV圖像讀寫與后處理安裝 POTpip install POT如果你已經(jīng)安裝了 PyTorch可以用下面命令驗(yàn)證 POT 是否安裝成功python -c import ot; print(ot.__version__)4.2 計(jì)算最優(yōu)傳輸映射下面的函數(shù)接收兩個(gè)二值掩膜返回一個(gè)從源掩膜到目標(biāo)掩膜的傳輸計(jì)劃import numpy as np import ot def compute_ot_map(mask_src, mask_tgt, eps1e-6): 計(jì)算從源掩膜到目標(biāo)掩膜的最優(yōu)傳輸計(jì)劃。 參數(shù): mask_src (np.ndarray): 源病灶掩膜形狀 (H, W)像素值為 0/1 mask_tgt (np.ndarray): 目標(biāo)病灶掩膜形狀 (H, W)像素值為 0/1 eps: 防止除零的極小值 返回: G (np.ndarray): 傳輸計(jì)劃形狀 (H, W, H, W) G[i, j, a, b] 表示從源像素 (i, j) 搬運(yùn)到目標(biāo)像素 (a, b) 的質(zhì)量 h, w mask_src.shape # 1. 將掩膜展平成概率分布 mu_s mask_src.reshape(-1).astype(np.float64) mu_t mask_tgt.reshape(-1).astype(np.float64) mu_s / (mu_s.sum() eps) mu_t / (mu_t.sum() eps) # 2. 生成像素坐標(biāo)網(wǎng)格 yy, xx np.mgrid[0:h, 0:w] coords np.stack([xx.ravel(), yy.ravel()], axis1).astype(np.float64) # 3. 計(jì)算代價(jià)矩陣歐氏距離的平方 M ot.dist(coords, coords, metricsqeuclidean) M / (M.max() eps) # 4. 使用 EMD 求解最優(yōu)傳輸計(jì)劃 G ot.emd(mu_s, mu_t, M) return G.reshape(h, w, h, w)這段代碼最核心的是兩個(gè)部分ot.dist(coords, coords, metricsqeuclidean)計(jì)算所有源像素與所有目標(biāo)像素之間的代價(jià)矩陣。ot.emd(mu_s, mu_t, M)用線性規(guī)劃求解最優(yōu)傳輸計(jì)劃。代價(jià)矩陣的標(biāo)準(zhǔn)通常會(huì)顯著影響數(shù)值穩(wěn)定性所以我對(duì) M 做了歸一化處理讓最大值等于 1。實(shí)際使用中如果你的病灶尺寸較大比如 128×128直接展開像素會(huì)得到 16384×16384 的代價(jià)矩陣內(nèi)存開銷極高建議先在連通域級(jí)別計(jì)算或者縮小到 32×32 再上采樣。4.3 計(jì)算 Wasserstein 重心POT 提供了ot.bregman.convolutional_barycenter2d可以直接對(duì)多張二維圖像計(jì)算熵正則化的 Wasserstein 重心def wasserstein_barycenter(masks, reg1e-2, num_iters100): 計(jì)算多張掩膜的 Wasserstein 重心。 參數(shù): masks (np.ndarray): 形狀 (N, H, W) 的二值掩膜數(shù)組 reg: 熵正則化系數(shù)越大速度越快但形狀越模糊 num_iters: Sinkhorn 迭代次數(shù) 返回: bary (np.ndarray): 形狀 (H, W) 的密度圖值在 0~1 之間 n, h, w masks.shape # 歸一化為概率分布 A np.stack([ m.astype(np.float64) / (m.sum() 1e-6) for m in masks ], axis0) bary ot.bregman.convolutional_barycenter2d( A, reg, numItermaxnum_iters ) return bary需要注意convolutional_barycenter2d使用的是卷積 Sinkhorn 算法計(jì)算效率比普通 Sinkhorn 高很多適合圖像這種二維網(wǎng)格數(shù)據(jù)。reg是一個(gè)關(guān)鍵參數(shù)reg 太小形狀更銳利但迭代收斂慢。reg 太大計(jì)算快但得到的重心會(huì)偏模糊。由于我們最終需要合成清晰的病灶掩膜通常會(huì)在得到重心密度圖后用閾值比如 0.5進(jìn)行二值化或者再做一次連通域篩選去掉零碎噪聲。4.4 紋理搬運(yùn)與合成病灶傳輸計(jì)劃的作用是把源病灶紋理搬運(yùn)到目標(biāo)位置。一個(gè)樸素的實(shí)現(xiàn)如下def transport_texture(image_src, G): 根據(jù)傳輸計(jì)劃 G 將源圖像紋理搬運(yùn)到目標(biāo)坐標(biāo)位置。 參數(shù): image_src (np.ndarray): 源病灶圖像形狀 (H, W, C) 或 (H, W) G (np.ndarray): 傳輸計(jì)劃形狀 (H, W, H, W) 返回: warped (np.ndarray): 搬運(yùn)后的紋理圖形狀與 image_src 一致 h, w image_src.shape[:2] channel_dim image_src.ndim warped np.zeros_like(image_src, dtypenp.float64) # 對(duì)每個(gè)目標(biāo)像素 (a, b) 加權(quán)聚合源像素的值 for a in range(h): for b in range(w): weights G[:, :, a, b] total weights.sum() if total 1e-12: continue weights weights / total if channel_dim 2: warped[a, b] (image_src * weights).sum() else: for c in range(image_src.shape[2]): warped[a, b, c] (image_src[:, :, c] * weights).sum() return warped這個(gè)雙重循環(huán)在理論上是正確的但效率很低。實(shí)際工程中可以先把傳輸計(jì)劃轉(zhuǎn)化為一個(gè)稀疏的“源坐標(biāo)重映射表”再用scipy.ndimage.map_coordinates做一次插值速度會(huì)快得多。下面把整個(gè) OTLesMix 合成流程串起來(lái)def otlesmix_synthesize(image_src, mask_src, image_tgt, mask_tgt, lam0.5): 合成一個(gè)位于源病灶與目標(biāo)病灶之間的新病灶。 參數(shù): image_src: 源病灶圖像塊形狀 (H, W, 3) mask_src: 源病灶掩膜形狀 (H, W)0/1 image_tgt: 目標(biāo)病灶圖像塊形狀 (H, W, 3) mask_tgt: 目標(biāo)病灶掩膜形狀 (H, W)0/1 lam: Wasserstein 重心的權(quán)重越大越接近目標(biāo)病灶 返回: synth_image: 合成病灶圖像塊 synth_mask: 合成病灶掩膜 # 1. 計(jì)算源到目標(biāo)的最優(yōu)傳輸計(jì)劃 G compute_ot_map(mask_src, mask_tgt) # 2. 將源紋理搬運(yùn)到目標(biāo)病灶位置 warped_texture transport_texture(image_src, G) # 3. 計(jì)算兩個(gè)掩膜的 Wasserstein 重心 bary wasserstein_barycenter( np.stack([mask_src, mask_tgt], axis0), reg1e-2 ) # 4. 根據(jù) lam 生成中間形狀掩膜 synth_mask bary (1 - lam) # 5. 背景部分保留目標(biāo)圖像病灶區(qū)域使用搬運(yùn)后的紋理 synth_image image_tgt.copy() synth_image[synth_mask] warped_texture[synth_mask] return synth_image, synth_mask.astype(np.float32)這段代碼中l(wèi)am既是重心權(quán)重也充當(dāng)了掩膜閾值的調(diào)節(jié)。實(shí)際使用中可以拆分成兩個(gè)獨(dú)立參數(shù)便于控制“形狀接近程度”和“掩膜閾值”。5. 從想法到實(shí)驗(yàn)如何驗(yàn)證 OTLesMix 的有效性設(shè)計(jì)好合成方法之后最關(guān)鍵的驗(yàn)證工作是把合成樣本真正用到分割模型訓(xùn)練中并和 baseline 對(duì)比。這里的實(shí)驗(yàn)協(xié)議值得認(rèn)真設(shè)計(jì)。5.1 評(píng)測(cè)任務(wù)選擇OTLesMix 適合作為病灶分割任務(wù)的數(shù)據(jù)增強(qiáng)模塊。常見的驗(yàn)證數(shù)據(jù)集包括BraTS2021多模態(tài)腦膠質(zhì)瘤 MRI 數(shù)據(jù)集包含 T1、T1ce、T2、FLAIR 四個(gè)模態(tài)標(biāo)注有壞死、水腫、增強(qiáng)腫瘤等結(jié)構(gòu)。LiTS肝臟及肝臟腫瘤 CT 數(shù)據(jù)集常用于腹部病灶分割。MSDMedical Segmentation Decathlon包含多種器官與病灶分割任務(wù)如胰腺、肝臟、結(jié)腸癌等。選擇數(shù)據(jù)集時(shí)最好選擇病灶形狀差異大、位置分布廣的數(shù)據(jù)這樣才能體現(xiàn) OTLesMix 的優(yōu)勢(shì)。如果數(shù)據(jù)集里全是圓心近似、大小相近的小病灶OT 方法和普通旋轉(zhuǎn)縮放的區(qū)別就不明顯。5.2 對(duì)比基線設(shè)置至少需要設(shè)置以下幾組對(duì)比方法說(shuō)明無(wú)增強(qiáng)直接用原始數(shù)據(jù)訓(xùn)練幾何增強(qiáng)旋轉(zhuǎn)、翻轉(zhuǎn)、縮放、彈性形變Mixup / CutMix通用混合增強(qiáng)作為像素混合基線OTLesMix本文討論的最優(yōu)傳輸合成方法評(píng)價(jià)指標(biāo)建議同時(shí)看全局指標(biāo)和困難樣本指標(biāo)Dice 系數(shù)最常用的區(qū)域重疊指標(biāo)。IoU交并比對(duì)邊界誤差更敏感。HD9595% 豪斯多夫距離反映邊界最大偏差對(duì)病灶邊緣質(zhì)量敏感。罕見形狀子集指標(biāo)把測(cè)試集中形狀最不規(guī)則、位置最邊緣的樣本單獨(dú)統(tǒng)計(jì)更容易看出 OTLesMix 帶來(lái)的提升。5.3 訓(xùn)練流程與超參數(shù)建議在實(shí)際訓(xùn)練中OTLesMix 通常作為一個(gè)在線增強(qiáng)器使用。推薦流程如下離線抽取訓(xùn)練集中所有病灶 patch 和對(duì)應(yīng)掩膜構(gòu)建病灶庫(kù)。每個(gè)訓(xùn)練 step 以一定概率比如 0.3~0.5觸發(fā) OTLesMix 增強(qiáng)。每次觸發(fā)時(shí)從病灶庫(kù)隨機(jī)抽取兩個(gè)病灶隨機(jī)選擇 λ ∈ [0.1, 0.9]。用 OTLesMix 生成合成 patch粘貼到當(dāng)前訓(xùn)練圖像中。使用合成 patch 的掩膜作為對(duì)應(yīng)位置的真實(shí)標(biāo)簽參與損失計(jì)算。這里有一個(gè)值得注意的細(xì)節(jié)合成病灶粘貼到訓(xùn)練圖時(shí)需要避免粘貼到不合理的解剖位置。最簡(jiǎn)單的方式是保持源病灶與目標(biāo)病灶的相對(duì)解剖位置關(guān)系或者限制粘貼位置在特定器官區(qū)域內(nèi)。6. 常見理解誤區(qū)與排查建議在理解和復(fù)現(xiàn) OTLesMix 的過(guò)程中下面幾個(gè)誤區(qū)比較常見我整理成表格方便排查。誤區(qū)產(chǎn)生原因正確理解與解決思路把 Wasserstein 重心當(dāng)作逐像素線性插值沒有理解分布的“質(zhì)量流動(dòng)”概念線性插值在歐氏空間逐點(diǎn)平均容易產(chǎn)生模糊Wasserstein 重心是概率分布之間的幾何平均能保留結(jié)構(gòu)直接把 OTLesMix 當(dāng) GAN 用期望生成全新語(yǔ)義混淆“融合已有樣本”和“從潛空間采樣生成”O(jiān)TLesMix 是在真實(shí)樣本之間插值語(yǔ)義由真實(shí)樣本約束不會(huì)憑空創(chuàng)造病灶代價(jià)矩陣用像素歐氏距離但不歸一化EMD 求解數(shù)值不穩(wěn)定建議先對(duì)代價(jià)矩陣做歸一化處理或改用熵正則化 Sinkhorn 求解只生成病灶但不檢查視覺合理性實(shí)驗(yàn)流程不夠嚴(yán)謹(jǐn)至少要做人工抽樣檢查保證紋理連續(xù)、邊界不出現(xiàn)明顯偽影在整張圖上計(jì)算 OT內(nèi)存溢出圖像尺寸太大先裁剪病灶區(qū)域到小 patch 上計(jì)算合成后再貼回原圖認(rèn)為 λ 越極端多樣性越好對(duì)插值系數(shù)作用理解不深λ 接近 0 或 1 時(shí)合成結(jié)果接近真實(shí)樣本多樣性有限0.3~0.7 區(qū)間更值得探索如果你遇到合成結(jié)果中出現(xiàn)大量空洞或者噪聲點(diǎn)優(yōu)先檢查掩膜歸一化是否遺漏以及reg正則化系數(shù)是否過(guò)大。reg過(guò)大會(huì)讓重心變得過(guò)于平滑二值化后很容易出現(xiàn)零碎偽影。7. 工程化落地的幾條建議把 OTLesMix 從論文實(shí)驗(yàn)遷移到實(shí)際項(xiàng)目中有幾個(gè)工程層面的問(wèn)題值得提前規(guī)劃。7.1 病灶庫(kù)的構(gòu)建與管理病灶庫(kù)的質(zhì)量直接決定合成樣本質(zhì)量。建議在訓(xùn)練流程開始前先對(duì)訓(xùn)練集做一次離線分析篩選出面積適中、邊界清晰的病灶。太小的病灶比如只有十幾個(gè)像素在 OT 計(jì)算中噪聲非常大邊界模糊的病灶掩膜本身就不準(zhǔn)合成出來(lái)的樣本也會(huì)帶偏模型。病灶庫(kù)可以采用統(tǒng)一尺寸的 patch 存儲(chǔ)例如統(tǒng)一縮放到 64×64 或 128×128。為了避免病灶形態(tài)被縮放扭曲記錄每個(gè) patch 的原始尺寸和縮放比例合成后再變換回原圖尺度。7.2 在線增強(qiáng) vs 離線增強(qiáng)OTLesMix 的計(jì)算開銷主要集中在 EMD 求解上在線增強(qiáng)時(shí)如果每個(gè) batch 都重新計(jì)算會(huì)拖慢訓(xùn)練速度。實(shí)際工程中推薦“離線預(yù)計(jì)算 在線隨機(jī)組合”的方式離線階段從病灶庫(kù)中挑選有代表性的配對(duì)組合預(yù)計(jì)算傳輸計(jì)劃和重心形狀保存到磁盤。在線階段訓(xùn)練時(shí)直接讀取預(yù)計(jì)算的組合結(jié)果隨機(jī)選擇 λ 進(jìn)行紋理搬運(yùn)和粘貼將單次增強(qiáng)開銷控制在毫秒級(jí)。這種方式犧牲了一部分隨機(jī)性但大幅提高了訓(xùn)練效率更適合大規(guī)模數(shù)據(jù)場(chǎng)景。7.3 質(zhì)量控制與醫(yī)學(xué)合規(guī)合成病灶生成得再逼真也不能繞過(guò)醫(yī)學(xué)數(shù)據(jù)的合規(guī)要求。使用真實(shí)病人數(shù)據(jù)構(gòu)建病灶庫(kù)前必須確認(rèn)數(shù)據(jù)的合法授權(quán)、匿名化處理和倫理審批。合成樣本不應(yīng)包含可識(shí)別患者身份的信息也不能當(dāng)成真實(shí)影像用于任何臨床診斷目的。此外如果項(xiàng)目要上線到真實(shí)醫(yī)療場(chǎng)景建議邀請(qǐng)影像科醫(yī)生對(duì)合成樣本進(jìn)行抽樣評(píng)估確認(rèn)病灶的形態(tài)、邊緣和紋理在病理上合理。這一點(diǎn)是最容易被算法團(tuán)隊(duì)忽視但又是最重要的環(huán)節(jié)。7.4 增強(qiáng)策略的動(dòng)態(tài)調(diào)整不建議在訓(xùn)練全程使用固定的 OTLesMix 觸發(fā)概率。更合理的做法是課程式增強(qiáng)訓(xùn)練早期模型還在學(xué)習(xí)基礎(chǔ)特征以常規(guī)幾何增強(qiáng)為主訓(xùn)練中期逐步引入 OTLesMix提高樣本多樣性訓(xùn)練后期收斂階段再降低觸發(fā)概率避免模型過(guò)于依賴合成樣本而產(chǎn)生分布偏移。你可以用一個(gè)簡(jiǎn)單的 epoch 調(diào)度函數(shù)來(lái)控制觸發(fā)概率def get_otlesmix_prob(epoch, max_epochs, p_max0.5): 隨著訓(xùn)練進(jìn)行先提升再降低 OTLesMix 使用概率。 前半程線性上升到 p_max后半程線性下降。 half max_epochs / 2 if epoch half: return p_max * (epoch / half) else: return p_max * (1 - (epoch - half) / half)7.5 與其它增強(qiáng)方法的組合OTLesMix 并不排斥傳統(tǒng)增強(qiáng)方法。推薦在 OTLesMix 合成結(jié)束后再疊加隨機(jī)旋轉(zhuǎn)、小幅度縮放和亮度擾動(dòng)進(jìn)一步增加多樣性。但要注意控制組合強(qiáng)度避免病灶紋理被過(guò)度擾動(dòng)而失真。組合增強(qiáng)的順序也影響效果。通常建議先做空間變換旋轉(zhuǎn)、翻轉(zhuǎn)再做強(qiáng)度變換亮度、對(duì)比度最后再做隨機(jī)遮擋類增強(qiáng)。對(duì)醫(yī)學(xué)圖像而言空間變換要保持解剖關(guān)系所以旋轉(zhuǎn)角度不宜過(guò)大比如限制在 ±15 度以內(nèi)。8. 總結(jié)與下一步學(xué)習(xí)方向OTLesMix 的價(jià)值在于提供了一種真正可解釋、可控制的病灶合成思路。它把病灶分割中的數(shù)據(jù)稀缺問(wèn)題轉(zhuǎn)化為經(jīng)典的最優(yōu)傳輸數(shù)學(xué)問(wèn)題用 Wasserstein 重心生成多樣形狀用最優(yōu)傳輸映射搬運(yùn)紋理和位置。這種方法不需要訓(xùn)練額外的生成模型也不會(huì)出現(xiàn) GAN 常見的訓(xùn)練不穩(wěn)定問(wèn)題。如果你想把 OTLesMix 用在自己的任務(wù)里我的建議是從小規(guī)模開始先拿一個(gè)小數(shù)據(jù)集只做兩個(gè)病灶之間的 OT 混合看看生成的 patch 在視覺上是否合理確認(rèn)紋理和形狀都自然之后再把它接入訓(xùn)練管線。最優(yōu)傳輸?shù)臄?shù)學(xué)看起來(lái)硬核但它最終解決的問(wèn)題非常直觀——把一個(gè)分布搬到另一個(gè)分布。理解了這句話就理解了 OTLesMix 的一半。接下來(lái)你可以繼續(xù)深入的方向包括Sinkhorn 算法與熵正則化的原理、切片 Wasserstein 距離在高維問(wèn)題中的應(yīng)用、最優(yōu)傳輸在域適應(yīng)中的用法以及如何用卷積 Sinkhorn 加速圖像級(jí)重心計(jì)算。如果這篇文章對(duì)你有幫助可以先收藏備用如果你在復(fù)現(xiàn)中遇到了具體報(bào)錯(cuò)或效果異常歡迎在評(píng)論區(qū)把你的處理步驟和現(xiàn)象發(fā)出來(lái)一起排查。