AAAI 2025丨2080Ti 也能 4K 圖像摳圖 !美圖&北交大提出超高解像度自然圖像摳圖算法 MEMatte

顯存友好的自然圖像摳圖框架用於在顯存受限的場景下(如商用顯卡與邊緣設備)進行高清圖像的精細摳圖。

日前,美圖影像研究院(MT Lab)聯合北京交通大學提出了超高解像度摳圖方法MEMatte (Memory Efficient),一個顯存友好的自然圖像摳圖框架用於在顯存受限的場景下(如商用顯卡與邊緣設備)進行高清圖像的精細摳圖。文章已被人工智能頂會 AAAI 2025 接收。

隨著攝影設備的性能升級,我們在日常採集到的圖像解像度已經能達到4K乃至8K,這種高解像度圖像有著更好的視覺質感與更豐富的細節內容。而自然圖像摳圖方法相比於一般的分割方法更善於區分細節充實的場景,因此適用於更複雜的下遊任務諸如圖像與影片的後期編輯。

當前Transformer (ViT) 為自然圖像摳圖方法帶來了顯著的性能提升,這種提升來自於全局注意力模塊對圖像內容的長程建模能力。然而,基於ViT的摳圖方法在計算資源受限的場景下無法處理高解像度圖像,因為ViT內部的全局注意力會產生次方倍的計算開銷。此外,摳圖任務的特性要求輸入能同時保留清晰圖像中的細節信息與完整圖像中的語義信息,因此無法採用降采樣或圖像切片等手段來降低模型輸入的圖像尺寸。這些限制條件要求我們在維持圖像輸入大小的前提下降低全局注意力對計算資源的消耗。

過往的解決方案包括令牌剪枝與令牌融合,這些方法都需要在模型的不同層以固定比例逐漸丟棄冗餘令牌。但是這類方法存在兩種缺陷:

(1)丟棄令牌會導致摳圖效果顯著變差,因為被丟棄的令牌中存在著需要被保留的細節信息;

(2)預設固定的丟棄比例無法適應不同的輸入,對於複雜樣例丟棄過多信息會導致精度下降,而對於簡單樣例則依然存在計算冗餘。

為瞭解決上述問題,研究團隊提出的MEMatte採用了雙分支的令牌路由設計。該設計的核心思想是將包含語義信息的令牌送入全局注意力模塊進行處理,並且將其餘令牌送入精心設計的輕量化令牌提煉模塊 (Light-weight Token Refinement Module, LTRM) 進行處理。令牌的分流是通過動態路由機制 (Batch-constrained Adaptive Token Routing, BATR)進行篩選,在推理時無需預設固定的比例,因此能夠輸入自適應地進行調整。這一設計顯著降低了模型的計算開銷,並在Nvidia GeForce 2080Ti 商用 GPU上實現了4K解像度圖像摳圖。此外,研究團隊還開源了超高解像度自然圖像摳圖數據集UHR-395 (Ultra High Resolution dataset),用於高解像度模型的訓練與評估。

MEMatte方法與已有方法在不同解像度的顯存消耗對比對比MEMatte方法與已有方法在不同解像度的顯存消耗對比對比
UHR-395數據集與其它開源摳圖數據集的輸入下UHR-395有著最高的平均解像度UHR-395數據集與其它開源摳圖數據集的輸入下UHR-395有著最高的平均解像度

1

方法

在推理過程中,MEMatte的ViT編碼器內部的每一個全局注意力模塊前都存在一個路由器 (Router)。這個路由器通過局部-全局策略為全局注意力進行令牌重要性評估。該策略對每層的圖像輸入進行重新映射以及通道拆分後的全局池化操作來兼顧局部與全局信息。此後通過LogSoftmax (LS)操作形成令牌的分流結果。

Router 將分流到全局注意力之外的令牌送入 LTRM 進行處理。LTRM 由數個輕量化組建構成,其中包括兩組映射層。此外,深度卷積 (Depth-Wise Convolution, DWC) 層處理局部空間信息,高效通道注意力層 (Efficient Channel Attention, ECA) 處理全局特徵信息。

在訓練MEMatte時,研究團隊採用 BATR 的訓練機制,搭配蒸餾學習策略與常規摳圖訓練損失,實現快速且平穩的模型訓練,在訓練過程中,假定從數據分佈中均勻地采樣,則每次的壓縮比例應當為一個定值。基於這項假設, BATR 為模型預設了分流比例

為如下形式:

來約束模型預測的批次平均分流比例,定義批次平均分流比例

通過約束模型的預測值與預設值的差異,訓練出足夠有效的 Router 以自適應的處理各種圖像。

MEMatte的訓練與推理流程MEMatte的訓練與推理流程

2

試驗結果

實驗結果表明,MEMatte在現有的基準以及UHR-395的高解像度測試集上都領先於過去的方法,並在域外測試集上展現了良好的泛化性能。此外,在摳圖測試集Composition-1K上,MEMatte相比於其baseline節省了約88%的顯存開銷,並降低了約50%的推理時間。

MEMatte與過往方法在通用摳圖基準上的比較MEMatte與過往方法在通用摳圖基準上的比較
MEMatte與過往方法在UHR-395高解像度測試集上的比較MEMatte與過往方法在UHR-395高解像度測試集上的比較
在高解像度圖像上的摳圖效果在高解像度圖像上的摳圖效果

研究圖案對 Router 分流結果的可視化結果展示了MEMatte在不同的全局注意力層對令牌的選擇傾向。MeMatte使中間的自注意力層主要關注物體邊緣與細區域的複雜區域,同時避免了令牌通過ViT淺層的自注意力操作,這大幅提高了輸入自注意力層的令牌稀疏性。

對Router分流結果的可視化,其中灰色區域是分流向 LTRM 的令牌對Router分流結果的可視化,其中灰色區域是分流向 LTRM 的令牌

3

超高解像度自然圖像摳圖數據集

目前公開的摳圖數據集難以作為高解像度摳圖的基準,因為它們無法同時兼顧類別多樣性和高解像度。為瞭解決這一問題,研究團隊提出了一個超高解像度自然摳圖數據集UHR-395。該數據集包含11個類別,共395個前景物體,劃分為訓練集355張以及測試集40張,平均解像度高達4872×6017。

前景類別分佈前景類別分佈
圖像解像度分佈圖像解像度分佈

為確保標註數據的高質量,每個前景物體的標註結果都通過多輪嚴格的審核機制,並由摳圖領域專家進行最後確認,不合格標註結果均被返回進行重新修正,最終獲得高質量的自然圖像摳圖標註數據。

標註結果展示標註結果展示

4

結語

超高解像度的自然圖像摳圖方法對實現複雜圖像編輯與影片後期處理等下遊任務而言至關重要,在本研究中,來自美圖影像研究院(MT Lab)和北京交通大學的研究人員共同提出了 MEMatte 和超高解像度數據集 UHR-395。與現有的摳圖方法相比,採用MEMatte 能夠顯著降低推理過程中的內存使用量和延遲,在高解像度和真實世界數據集上實現了最先進的性能,突出了令牌路由機制 BATR和 輕量化令牌細化模塊 LTRM 的穩健性,為美圖的圖像與影片處理能力提供高性能且靈活的技術支撐,為用戶帶來高效優質使用體驗。此外, MEMatte 也有望促進高解像度摳圖的進一步研究。