AAAI2025丨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方法與已有方法在不同解析度的視訊記憶體消耗對比對比
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的訓練與推理流程
2
試驗結果
實驗結果表明,MEMatte在現有的基準以及UHR-395的高解析度測試集上都領先於過去的方法,並在域外測試集上展現了良好的泛化效能。此外,在摳圖測試集Composition-1K上,MEMatte相比於其baseline節省了約88%的視訊記憶體開銷,並降低了約50%的推理時間。
MEMatte與過往方法在通用摳圖基準上的比較
MEMatte與過往方法在UHR-395高解析度測試集上的比較
在高解析度影像上的摳圖效果
研究圖案對 Router 分流結果的視覺化結果展示了MEMatte在不同的全域性注意力層對令牌的選擇傾向。MeMatte使中間的自注意力層主要關注物體邊緣與細區域的複雜區域,同時避免了令牌透過ViT淺層的自注意力操作,這大幅提高了輸入自注意力層的令牌稀疏性。
對Router分流結果的視覺化,其中灰色區域是分流向 LTRM 的令牌
3
超高解析度自然影像摳圖資料集
目前公開的摳圖資料集難以作為高解析度摳圖的基準,因為它們無法同時兼顧類別多樣性和高解析度。為了解決這一問題,研究團隊提出了一個超高解析度自然摳圖資料集UHR-395。該資料集包含11個類別,共395個前景物體,劃分為訓練集355張以及測試集40張,平均解析度高達4872×6017。
前景類別分佈
影像解析度分佈
為確保標註資料的高質量,每個前景物體的標註結果都透過多輪嚴格的稽核機制,並由摳圖領域專家進行最後確認,不合格標註結果均被返回進行重新修正,最終獲得高質量的自然影像摳圖標註資料。
標註結果展示
4
結語
超高解析度的自然影像摳圖方法對實現複雜影像編輯與影片後期處理等下游任務而言至關重要,在本研究中,來自美圖影像研究院(MT Lab)和北京交通大學的研究人員共同提出了 MEMatte 和超高解析度資料集 UHR-395。與現有的摳圖方法相比,採用MEMatte 能夠顯著降低推理過程中的記憶體使用量和延遲,在高解析度和真實世界資料集上實現了最先進的效能,突出了令牌路由機制 BATR和 輕量化令牌細化模組 LTRM 的穩健性,為美圖的影像與影片處理能力提供高效能且靈活的技術支撐,為使用者帶來高效優質使用體驗。此外, MEMatte 也有望促進高解析度摳圖的進一步研究。
更多內容,點選下方關注:

未經「AI科技評論」授權,嚴禁以任何方式在網頁、論壇、社群進行轉載!
公眾號轉載請先在「AI科技評論」後臺留言取得授權,轉載時需標註來源並插入本公眾號名片。


相關文章