邁向分割的大一統(tǒng)!OMG-Seg:一個模型搞定所有分割任務(wù)
本文經(jīng)自動駕駛之心公眾號授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
寫在前面&筆者個人思考
圖像分割已經(jīng)從單任務(wù)分割走到了語義分割、實例分割、全景分割三種分割任務(wù)的統(tǒng)一;大模型以及多模態(tài)的發(fā)展又帶來了文本和圖像統(tǒng)一,使得跨模態(tài)端到端成為可能;追求更高級、更全面任務(wù)似乎成了CV/NLP領(lǐng)域的共識。聽起來很玄乎,對于本文的統(tǒng)一所有分割來說,做法也很好理解,本質(zhì)上就是在Mask2former的基礎(chǔ)上增加了SAM的提示來支持prompt驅(qū)動以及交互式分割,同時增加CLIP encoder來支持開放域,給query加ID使其同時支持圖像和視頻的分割;
PS:不得不說,Mask2former還是一個很強的基線,當前SOTA的大一統(tǒng)分割模型Oneformer是在這個基礎(chǔ)上改進的,本文的超大一統(tǒng)方案也基于這個base進行升級;
論文的主要思路
本文統(tǒng)一解決了各種分割任務(wù),提出的OMG-Seg,即一種模型,足夠出色,能夠高效而有效地處理所有分割任務(wù),包括圖像語義分割、實例分割、全景分割,以及它們的視頻對應(yīng)任務(wù)、開放詞匯設(shè)置、由提示驅(qū)動的交互分割(如SAM)和視頻目標分割。據(jù)本文所知,這是第一個能夠處理所有這些任務(wù)并取得令人滿意性能的模型。本文展示了 OMG-Seg,這是一個基于Transformer的編碼器-解碼器架構(gòu),具有任務(wù)特定的查詢和輸出,可以支持十多種不同的分割任務(wù),并顯著減少跨各種任務(wù)和數(shù)據(jù)集的計算和參數(shù)開銷。本文在共同訓(xùn)練期間嚴格評估了任務(wù)間的影響和相關(guān)性。代碼和模型可在 https://github.com/lxtGH/OMG-Seg 獲取。
主要貢獻
引入OMG-Seg,一個模型應(yīng)對所有分割任務(wù)。
結(jié)構(gòu)上:
- 先前的模型通常使用共享的視覺骨干,有幾個特定任務(wù)的分支,但OMG-Seg采用了共享的編碼器-解碼器架構(gòu)。
- 將所有任務(wù)的輸出統(tǒng)一為一個統(tǒng)一的查詢表示。一個查詢可以表示一個掩模標簽/圖像或掩模/唯一的ID/視覺提示。然后,采用一個共享的解碼器處理所有類型的查詢及其特征。這個設(shè)置促進了通用分割訓(xùn)練和推理的統(tǒng)一,并實現(xiàn)了跨任務(wù)的廣泛參數(shù)共享。
效果上:
- 通過在合并的圖像和視頻數(shù)據(jù)集上進行共同訓(xùn)練,OMG-Seg能夠處理多達十種不同的分割任務(wù),涵蓋不同的數(shù)據(jù)集。OMG-Seg是首個在圖像、視頻、開放詞匯和交互分割四種設(shè)置上在八個不同的數(shù)據(jù)集上取得了不錯的結(jié)果,包括COCO、ADE20k、VIPSeg、Youtube-VIS-2019、Youtube-VIS-2021和DAVIS-1。
相關(guān)工作:
通用圖像/視頻分割 視覺Transformer的出現(xiàn)引發(fā)了通用分割領(lǐng)域的創(chuàng)新浪潮。最近的研究開發(fā)了基于端到端集合預(yù)測方法的掩模分類架構(gòu),在圖像和視頻分割任務(wù)中表現(xiàn)出色,超過了專業(yè)模型。盡管取得了這些進展,大多數(shù)現(xiàn)有方法仍然依賴于不同的模型來處理不同的分割任務(wù)和數(shù)據(jù)集。最近,有一種趨勢是在各種數(shù)據(jù)集和任務(wù)之間訓(xùn)練單一模型,以獲得參數(shù)共享的好處。例如,OneFormer在一個單一模型中集成了三個圖像分割任務(wù),而UNINEXT專注于統(tǒng)一實例級任務(wù)。類似地,TarVIS使用目標提示組合各種視頻分割任務(wù)。然而,這些現(xiàn)有工作中沒有一個徹底研究過圖像、視頻和由提示驅(qū)動的數(shù)據(jù)的聯(lián)合訓(xùn)練,以構(gòu)建一個全面的分割模型。本文的工作是朝著這個方向的首次嘗試,拓展了跨這些領(lǐng)域的共同訓(xùn)練的潛力。
視覺基礎(chǔ)模型 視覺基礎(chǔ)模型領(lǐng)域的最新研究展示了優(yōu)化技術(shù)的多樣化,涵蓋了各種學(xué)習(xí)范式。這些包括僅視覺預(yù)訓(xùn)練策略、聯(lián)合視覺-語言預(yù)訓(xùn)練方法以及包含視覺提示的多模態(tài)框架。一個值得注意的例子是SAM,它展示了廣泛培訓(xùn)在實現(xiàn)通用分割方面的泛化性和可擴展性。在此基礎(chǔ)上,Semantic-SAM通過添加語義標簽和增加細粒度級別的方式增強了SAM模型。然而,盡管它們的印象深刻,這些視覺基礎(chǔ)模型通常在視頻分割任務(wù)上表現(xiàn)不佳,需要在更動態(tài)的情境中進行進一步的優(yōu)化以實現(xiàn)最佳性能。
開放詞匯分割 這一系列的視覺分割研究旨在識別和分割超出有限封閉集視覺概念的新目標。許多研究利用視覺語言模型(VLMs)提供的可轉(zhuǎn)移表示,在訓(xùn)練過程中探索區(qū)域和文本表示之間的對齊。在推理階段,檢測器可以使用從VLMs派生的文本嵌入來識別新的類別。本文的模型遵循這一理念,以實現(xiàn)開放詞匯分割。具體來說,本文使用凍結(jié)的VLMs同時作為特征提取器和分類器。這種策略允許在開放詞匯設(shè)置中實現(xiàn)無縫過渡。
統(tǒng)一建模 Transformer架構(gòu)的適應(yīng)性使得在不同模態(tài)之間共享基本模塊成為可能。這種多功能性啟發(fā)了幾個研究倡議,使用一個通用的Transformer框架處理不同領(lǐng)域。值得注意的是,在視覺通用領(lǐng)域方面的努力主要集中在統(tǒng)一視覺領(lǐng)域內(nèi)不同任務(wù)。例如,Pix2Seq系列通過自回歸標記預(yù)測的方式實現(xiàn)任務(wù)統(tǒng)一。類似地,Unified-IO實現(xiàn)了一個序列到序列的流水線,將各種輸入和輸出轉(zhuǎn)換為離散的標記序列。此外,最近的進展探索了在上下文中學(xué)習(xí)作為結(jié)合各種視覺任務(wù)的手段。這些方法主要針對跨領(lǐng)域的任務(wù)統(tǒng)一。然而,統(tǒng)一分割模型與為特定目的構(gòu)建的分割模型之間的性能差距仍然是一個待解決的問題。
方法設(shè)計:
OMG-Seg是一個單一而多才多藝的模型,具有減少任務(wù)特定定制和最大參數(shù)共享的特點,可以支持各種分割任務(wù),使其成為一個適用于所有分割任務(wù)的模型。目標不是追求每個任務(wù)的最先進結(jié)果,而是增加一個可廣泛推廣的分割模型的建模能力,同時允許任務(wù)之間的廣泛知識共享。
統(tǒng)一任務(wù)表示(看看都支持哪些任務(wù)~~)
圖像分割 給定輸入圖像,圖像分割的目標是輸出一組掩模,其中表示二進制掩模的類別標簽,是掩模數(shù)量,是空間大小。根據(jù)類別標簽和掩模的范圍,本文報告了三種不同的分割任務(wù)的結(jié)果,包括語義分割(SS)、實例分割(IS)和全景分割(PS)。PS是SS和IS的統(tǒng)一,其中包含可計數(shù)的thing類別和不可計數(shù)的stuff類別。對于這三個任務(wù),本文采用掩模分類架構(gòu),其中每個掩模對應(yīng)一個語義標簽。
視頻分割 給定視頻剪輯輸入,其中表示幀數(shù),視頻分割的目標是獲取一個掩模tube ,其中是tube掩模數(shù)量,。表示tube掩模的類別標簽,表示每個tube掩模的實例ID。每個tube掩模可以分類到可計數(shù)的thing或不可計數(shù)的stuff類別,其中thing類別還分配了一個唯一的ID。對于stuff掩模,默認情況下跟蹤為零。當且任務(wù)只包含stuff類別,且所有thing類別沒有ID時,VPS變成了視頻語義分割(VSS)。如果重疊且只包含thing類別,且所有stuff類別被忽略,VPS變成了視頻實例分割(VIS)。視頻目標分割(VOS)的目標是跟蹤第一個框架的掩模,而不進行分類。受圖像分割的啟發(fā),本文還采用了tube掩模分類架構(gòu)來訓(xùn)練和沿時間維度鏈接短管。對于VOS,本文采用類別不可知的tube掩模訓(xùn)練,這類似于VPS和VIS。目標分割(VOS)的目標是跟蹤第一個框架的掩模,而不進行分類。受到圖像分割的啟發(fā),本文還采用了tube掩模分類架構(gòu)來訓(xùn)練和沿時間維度鏈接短管。對于VOS,本文采用了類別不可知的tube掩模訓(xùn)練,這與VPS和VIS類似。
交互式分割 SAM框架中的交互式分割同時使用圖像 和視覺提示 (例如點和框)作為輸入,輸出相應(yīng)的二進制圖像掩模 ,其中 是視覺提示的數(shù)量。每個視覺提示被編碼成一個目標查詢,這自然地可以成為解碼器的輸入。在本文的實驗中,本文使用共享的解碼器處理所有不同任務(wù)查詢。
開放詞匯和多數(shù)據(jù)集分割 任務(wù)的制定與先前的圖像和視頻分割相同。然而,這個設(shè)置超越了固定的標簽空間。特別是,它需要在各種數(shù)據(jù)集上進行開放集識別。同時,多數(shù)據(jù)集分割要求一個模型在不同數(shù)據(jù)集下分割更多的概念。作為一種常見做法,本文采用CLIP文本嵌入作為掩模分類器,避免了分類學(xué)沖突,并同時實現(xiàn)了開放集識別。因此,本文測量了視覺查詢特征與類別嵌入之間的距離,而不是學(xué)到的分類器。
所有的事物都在查詢中 如上所述,通過結(jié)合所有不同的設(shè)置,本文可以使用相同基于查詢的掩模分類框架表示所有輸出分割實體。特別是,一個目標查詢對應(yīng)一個掩模 ,標簽 和 ID 。根據(jù)不同的任務(wù)設(shè)置, 和 的格式和范圍是不同的。然而, 和 的格式和范圍是相似的。因此,將所有這些任務(wù)放入一個共享的編碼器和解碼器框架中,并共同訓(xùn)練一個模型用于所有分割任務(wù)是很自然的。
OMG-Seg架構(gòu)(具體是怎么做的~~)
OMG-Seg遵循Mask2Former的架構(gòu)設(shè)計。如圖所示,它包含一個骨干網(wǎng)絡(luò)、一個像素解碼器和一個掩模解碼器。區(qū)別在于以下幾個方面,包括凍結(jié)骨干的設(shè)計、合并的目標查詢(包括目標查詢和視覺提示),以及共享的多任務(wù)解碼器。根據(jù)不同的任務(wù)設(shè)置,解碼器輸出相應(yīng)的掩模和標簽。
使用VLM編碼器作為凍結(jié)骨干 為了實現(xiàn)開放詞匯的識別,在骨干部分,本文采用凍結(jié)的CLIP視覺模型作為特征提取器。本文使用ConvNeXt架構(gòu)來自O(shè)penCLIP。給定圖像/視頻輸入,VLM編碼器提取多尺度的凍結(jié)特征 ,以供進一步處理。
像素解碼器作為特征適配器 像素解碼器與Mask2Former相同,包含多階段的可變形注意力層。它將凍結(jié)特征 轉(zhuǎn)換為具有相同通道維度的融合特征 ,其中 是特征的層索引, 是最高分辨率的特征。
合并的目標查詢 如上所分析,每個目標查詢代表一種掩模輸出。然而,從功能的角度來看,圖像、視頻和交互模式代表不同的屬性。對于圖像,目標查詢側(cè)重于目標級別的定位和識別。對于視頻,目標查詢可能涉及時間一致性,例如同一目標跨越不同幀。對于交互式分割,目標查詢被迫定位特定區(qū)域。對于圖像和視頻輸入,本文采用目標查詢來表示圖像掩模或被跟蹤的tube掩模。因為兩者都需要語義標簽。本文將它們稱為語義查詢,。對于交互模式,按照SAM 的做法,本文采用提示編碼器將各種視覺提示編碼成與目標查詢相同形狀的查詢。本文將其稱為位置查詢,。因此,本文可以共享變壓器解碼器的相同接口。
共享的多任務(wù)解碼器 其主要操作是交叉注意力,接收合并的目標查詢 和 ) 和圖像/視頻特征 ,并輸出精煉的目標查詢。最終的掩模是通過精煉查詢和高分辨率特征 的點積獲得的。對于圖像語義級別的任務(wù),本文采用了與Mask2Former相同的過程。具體而言, 使用多尺度特征 進行掩碼交叉注意力 [18]。 是查詢,而 是鍵和值。然后,對精煉的查詢應(yīng)用多頭自注意力(MHSA)層。
對于視頻任務(wù),本文采用相同的交叉注意力設(shè)計。唯一的區(qū)別是金字塔特征 沿著時間維度與3D位置嵌入連接,這是先前的工作中的默認設(shè)置。合并的視頻特征和精煉的查詢用于預(yù)測tube掩模。
對于交互式分割,本文執(zhí)行相同的交叉注意力設(shè)計。然而,本文跳過自注意力,以避免在MHSA層中的掩模查詢之間發(fā)生交互,因為交互式分割只關(guān)心輸入視覺提示區(qū)域。在獲取精煉的目標查詢后,它經(jīng)過一個預(yù)測FFN,通常包括一個帶有ReLU激活層和線性投影層的3層感知器。所有的查詢都受到掩模分類損失和掩模預(yù)測損失的監(jiān)督。解碼過程以級聯(lián)的方式進行,每個特征金字塔有三個階段。
Training and Inference(多數(shù)據(jù)集&多任務(wù)放一起怎么訓(xùn)~~)
聯(lián)合圖像視頻數(shù)據(jù)集聯(lián)合訓(xùn)練 與首先在圖像數(shù)據(jù)集上進行預(yù)訓(xùn)練不同,本文的目標是僅聯(lián)合訓(xùn)練所有分割任務(wù)一次。所有訓(xùn)練目標都是三種不同情況下的一個實體標簽和掩模。實體可以是物體、雜項、類別無關(guān)的掩模及其相應(yīng)的標簽。注意,具有相同ID 的實例掩模形成tube掩模。在訓(xùn)練過程中,本文在預(yù)測的實體掩模和地面實體掩模之間應(yīng)用匈牙利匹配,以將目標查詢分配給視頻/圖像實體,然后監(jiān)督它們的預(yù)測掩模和分類。為了避免跨數(shù)據(jù)集的分類沖突,分類器被替換為CLIP文本嵌入。最終的損失函數(shù)表示為 。這里, 是掩模分類的交叉熵(CE)損失, 和 分別是分割的掩模交叉熵(CE)損失和Dice損失。
通用推斷 對于圖像分割,本文遵循Mask2Former的相同推斷過程。例如,對于 PS,本文根據(jù)排序后的分數(shù)合并物體和雜項。分數(shù)由 CLIP 文本嵌入生成。對于視頻分割任務(wù),對于 VIS 和 VPS,為了生成實例ID,本文遵循先前的工作,使用查詢匹配而不是引入額外的跟蹤組件。對于 VOS 任務(wù),本文在第一幀和其余幀之間采用掩模匹配。對于交互式分割任務(wù),本文遵循原始的 SAM,通過提供框和點提示來獲取二進制掩模。對于開放詞匯分割任務(wù),由于本文有一個凍結(jié)的 CLIP 編碼器,本文將掩模池化分數(shù)和學(xué)習(xí)分數(shù)與開放詞匯嵌入合并。
結(jié)合任務(wù)以獲取更多應(yīng)用 由于本文的模型可以執(zhí)行各種分割任務(wù),結(jié)合交互、開放詞匯和圖像/視頻分割任務(wù)可以導(dǎo)致一些新的應(yīng)用。例如,本文可以結(jié)合交互和視頻分割,實現(xiàn)靈活的基于提示驅(qū)動的視頻目標分割。或者本文可以將交互分割與開放詞匯設(shè)置結(jié)合起來,從而實現(xiàn)開放詞匯交互分割。
實驗結(jié)果:
數(shù)據(jù)集和評估指標 與常規(guī)設(shè)置不同,本文的目標是盡可能地探索在多個數(shù)據(jù)集上進行聯(lián)合訓(xùn)練。本文使用COCO panoptic、COCO-SAM、VIPSeg 和Youtube-VIS-2019(YT-VIS19)作為訓(xùn)練數(shù)據(jù)集。除了封閉集測試,本文還包括使用Youtube-VIS-2021、ADE-20k和DAVIS-2017數(shù)據(jù)集進行開放詞匯(OV)推理,它們的注釋在訓(xùn)練過程中未被使用。COCO-SAM 是通過使用真實邊界框,掩模中心點作為視覺提示來創(chuàng)建的。注釋是通過COCO全景掩模獲得的。此外,本文還在表3中包括了多數(shù)據(jù)集設(shè)置,以驗證本文的OMG-Seg的多數(shù)據(jù)集聯(lián)合訓(xùn)練的有效性。除了表2之外,本文添加了更多數(shù)據(jù)集,包括ADE-20k 和YT-VIS21 用于聯(lián)合訓(xùn)練。本文為每個數(shù)據(jù)集使用相應(yīng)的評估指標,包括 PQ、掩模 mAP、VPQ、tube mAP、J&F和 mIoU。
實驗細節(jié) 本文在MMDetection中實現(xiàn)本文的模型和所有其他基線。本文使用32個A100 GPU的分布式訓(xùn)練框架。每個小批次在每個GPU上有一張圖像。對于數(shù)據(jù)增強,本文采用大規(guī)模的抖動,與先前的工作一樣,以構(gòu)建強大的基線。對于每個表中的所有模型,本文采用相同的訓(xùn)練步驟。本文使用OpenCLIP來初始化骨干網(wǎng)絡(luò),并用其相應(yīng)的文本嵌入替換學(xué)到的分類器。對于圖像輸入,將它們視為偽視頻,通過將兩個圖像及其掩模連接成一個。本文采用不同的采樣率來平衡每個數(shù)據(jù)集的訓(xùn)練樣本。本文報告凍結(jié)和訓(xùn)練后的骨干網(wǎng)絡(luò)的結(jié)果供參考。
Main Results
雖然在COCO圖像分割基準上表現(xiàn)略有下降,但在VIPSeg數(shù)據(jù)集上取得了接近最先進水平的結(jié)果;開放詞匯視頻分割領(lǐng)域效果也不錯。
基于K-Net和 Mask2Former進行了多數(shù)據(jù)集聯(lián)合訓(xùn)練。聯(lián)合共同訓(xùn)練通常可以提高大多數(shù)視頻分割數(shù)據(jù)集的性能,從而實現(xiàn)了大幅度的模型參數(shù)減少(從1326M減少到221M)。這種改進在三個VPS和VIS數(shù)據(jù)集中是一致的,無論是否凍結(jié)了骨干網(wǎng)絡(luò)。
結(jié)論:
本文介紹了第一個用于圖像、視頻、開放詞匯和交互式分割的聯(lián)合共訓(xùn)練框架。本文的解決方案OMG-Seg是一個新穎而簡單的框架,它使用統(tǒng)一的查詢表示和共享的解碼器來處理多樣化的任務(wù)。首次有可能訓(xùn)練一個單一的分割模型,能夠在十個不同的任務(wù)上表現(xiàn)出與特定任務(wù)模型相媲美的性能。這種方法顯著減少了參數(shù)大小和在各種應(yīng)用程序的模型設(shè)計中需要專門工程的需求。預(yù)計本文高效且多功能的框架將成為多任務(wù)和多數(shù)據(jù)集分割的穩(wěn)健基線。
原文鏈接:https://mp.weixin.qq.com/s/Hz2gDOy5POxU_bFyJCxUDw