邁向分割的大一統(tǒng)！OMG-Seg：一個(gè)模型搞定所有分割任務(wù)

作者：Xiangtai Li 2024-01-24 09:24:19

今天為大家分享大一統(tǒng)分割算法—OMG-Seg！一個(gè)模型搞定所有分割任務(wù)，包括圖像語(yǔ)義分割、實(shí)例分割、全景分割，以及它們的視頻對(duì)應(yīng)任務(wù)、開(kāi)放詞匯設(shè)置、由提示驅(qū)動(dòng)的交互分割（如SAM）和視頻目標(biāo)分割！

本文經(jīng)自動(dòng)駕駛之心公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

寫(xiě)在前面&筆者個(gè)人思考

圖像分割已經(jīng)從單任務(wù)分割走到了語(yǔ)義分割、實(shí)例分割、全景分割三種分割任務(wù)的統(tǒng)一；大模型以及多模態(tài)的發(fā)展又帶來(lái)了文本和圖像統(tǒng)一，使得跨模態(tài)端到端成為可能；追求更高級(jí)、更全面任務(wù)似乎成了CV/NLP領(lǐng)域的共識(shí)。聽(tīng)起來(lái)很玄乎，對(duì)于本文的統(tǒng)一所有分割來(lái)說(shuō)，做法也很好理解，本質(zhì)上就是在Mask2former的基礎(chǔ)上增加了SAM的提示來(lái)支持prompt驅(qū)動(dòng)以及交互式分割，同時(shí)增加CLIP encoder來(lái)支持開(kāi)放域，給query加ID使其同時(shí)支持圖像和視頻的分割；

PS：不得不說(shuō)，Mask2former還是一個(gè)很強(qiáng)的基線(xiàn)，當(dāng)前SOTA的大一統(tǒng)分割模型Oneformer是在這個(gè)基礎(chǔ)上改進(jìn)的，本文的超大一統(tǒng)方案也基于這個(gè)base進(jìn)行升級(jí)；

論文的主要思路

本文統(tǒng)一解決了各種分割任務(wù)，提出的OMG-Seg，即一種模型，足夠出色，能夠高效而有效地處理所有分割任務(wù)，包括圖像語(yǔ)義分割、實(shí)例分割、全景分割，以及它們的視頻對(duì)應(yīng)任務(wù)、開(kāi)放詞匯設(shè)置、由提示驅(qū)動(dòng)的交互分割（如SAM）和視頻目標(biāo)分割。據(jù)本文所知，這是第一個(gè)能夠處理所有這些任務(wù)并取得令人滿(mǎn)意性能的模型。本文展示了 OMG-Seg，這是一個(gè)基于Transformer的編碼器-解碼器架構(gòu)，具有任務(wù)特定的查詢(xún)和輸出，可以支持十多種不同的分割任務(wù)，并顯著減少跨各種任務(wù)和數(shù)據(jù)集的計(jì)算和參數(shù)開(kāi)銷(xiāo)。本文在共同訓(xùn)練期間嚴(yán)格評(píng)估了任務(wù)間的影響和相關(guān)性。代碼和模型可在 https://github.com/lxtGH/OMG-Seg 獲取。

主要貢獻(xiàn)

引入OMG-Seg，一個(gè)模型應(yīng)對(duì)所有分割任務(wù)。

結(jié)構(gòu)上：

先前的模型通常使用共享的視覺(jué)骨干，有幾個(gè)特定任務(wù)的分支，但OMG-Seg采用了共享的編碼器-解碼器架構(gòu)。
將所有任務(wù)的輸出統(tǒng)一為一個(gè)統(tǒng)一的查詢(xún)表示。一個(gè)查詢(xún)可以表示一個(gè)掩模標(biāo)簽/圖像或掩模/唯一的ID/視覺(jué)提示。然后，采用一個(gè)共享的解碼器處理所有類(lèi)型的查詢(xún)及其特征。這個(gè)設(shè)置促進(jìn)了通用分割訓(xùn)練和推理的統(tǒng)一，并實(shí)現(xiàn)了跨任務(wù)的廣泛參數(shù)共享。

效果上：

通過(guò)在合并的圖像和視頻數(shù)據(jù)集上進(jìn)行共同訓(xùn)練，OMG-Seg能夠處理多達(dá)十種不同的分割任務(wù)，涵蓋不同的數(shù)據(jù)集。OMG-Seg是首個(gè)在圖像、視頻、開(kāi)放詞匯和交互分割四種設(shè)置上在八個(gè)不同的數(shù)據(jù)集上取得了不錯(cuò)的結(jié)果，包括COCO、ADE20k、VIPSeg、Youtube-VIS-2019、Youtube-VIS-2021和DAVIS-1。

方法設(shè)計(jì)：

OMG-Seg是一個(gè)單一而多才多藝的模型，具有減少任務(wù)特定定制和最大參數(shù)共享的特點(diǎn)，可以支持各種分割任務(wù)，使其成為一個(gè)適用于所有分割任務(wù)的模型。目標(biāo)不是追求每個(gè)任務(wù)的最先進(jìn)結(jié)果，而是增加一個(gè)可廣泛推廣的分割模型的建模能力，同時(shí)允許任務(wù)之間的廣泛知識(shí)共享。

統(tǒng)一任務(wù)表示（看看都支持哪些任務(wù)～～）

圖像分割 給定輸入圖像，圖像分割的目標(biāo)是輸出一組掩模，其中表示二進(jìn)制掩模的類(lèi)別標(biāo)簽，是掩模數(shù)量，是空間大小。根據(jù)類(lèi)別標(biāo)簽和掩模的范圍，本文報(bào)告了三種不同的分割任務(wù)的結(jié)果，包括語(yǔ)義分割（SS）、實(shí)例分割（IS）和全景分割（PS）。PS是SS和IS的統(tǒng)一，其中包含可計(jì)數(shù)的thing類(lèi)別和不可計(jì)數(shù)的stuff類(lèi)別。對(duì)于這三個(gè)任務(wù)，本文采用掩模分類(lèi)架構(gòu)，其中每個(gè)掩模對(duì)應(yīng)一個(gè)語(yǔ)義標(biāo)簽。

視頻分割 給定視頻剪輯輸入，其中表示幀數(shù)，視頻分割的目標(biāo)是獲取一個(gè)掩模tube ，其中是tube掩模數(shù)量，。表示tube掩模的類(lèi)別標(biāo)簽，表示每個(gè)tube掩模的實(shí)例ID。每個(gè)tube掩模可以分類(lèi)到可計(jì)數(shù)的thing或不可計(jì)數(shù)的stuff類(lèi)別，其中thing類(lèi)別還分配了一個(gè)唯一的ID。對(duì)于stuff掩模，默認(rèn)情況下跟蹤為零。當(dāng)且任務(wù)只包含stuff類(lèi)別，且所有thing類(lèi)別沒(méi)有ID時(shí)，VPS變成了視頻語(yǔ)義分割（VSS）。如果重疊且只包含thing類(lèi)別，且所有stuff類(lèi)別被忽略，VPS變成了視頻實(shí)例分割（VIS）。視頻目標(biāo)分割（VOS）的目標(biāo)是跟蹤第一個(gè)框架的掩模，而不進(jìn)行分類(lèi)。受圖像分割的啟發(fā)，本文還采用了tube掩模分類(lèi)架構(gòu)來(lái)訓(xùn)練和沿時(shí)間維度鏈接短管。對(duì)于VOS，本文采用類(lèi)別不可知的tube掩模訓(xùn)練，這類(lèi)似于VPS和VIS。目標(biāo)分割（VOS）的目標(biāo)是跟蹤第一個(gè)框架的掩模，而不進(jìn)行分類(lèi)。受到圖像分割的啟發(fā)，本文還采用了tube掩模分類(lèi)架構(gòu)來(lái)訓(xùn)練和沿時(shí)間維度鏈接短管。對(duì)于VOS，本文采用了類(lèi)別不可知的tube掩模訓(xùn)練，這與VPS和VIS類(lèi)似。

交互式分割 SAM框架中的交互式分割同時(shí)使用圖像和視覺(jué)提示（例如點(diǎn)和框）作為輸入，輸出相應(yīng)的二進(jìn)制圖像掩模，其中是視覺(jué)提示的數(shù)量。每個(gè)視覺(jué)提示被編碼成一個(gè)目標(biāo)查詢(xún)，這自然地可以成為解碼器的輸入。在本文的實(shí)驗(yàn)中，本文使用共享的解碼器處理所有不同任務(wù)查詢(xún)。

開(kāi)放詞匯和多數(shù)據(jù)集分割 任務(wù)的制定與先前的圖像和視頻分割相同。然而，這個(gè)設(shè)置超越了固定的標(biāo)簽空間。特別是，它需要在各種數(shù)據(jù)集上進(jìn)行開(kāi)放集識(shí)別。同時(shí)，多數(shù)據(jù)集分割要求一個(gè)模型在不同數(shù)據(jù)集下分割更多的概念。作為一種常見(jiàn)做法，本文采用CLIP文本嵌入作為掩模分類(lèi)器，避免了分類(lèi)學(xué)沖突，并同時(shí)實(shí)現(xiàn)了開(kāi)放集識(shí)別。因此，本文測(cè)量了視覺(jué)查詢(xún)特征與類(lèi)別嵌入之間的距離，而不是學(xué)到的分類(lèi)器。

所有的事物都在查詢(xún)中 如上所述，通過(guò)結(jié)合所有不同的設(shè)置，本文可以使用相同基于查詢(xún)的掩模分類(lèi)框架表示所有輸出分割實(shí)體。特別是，一個(gè)目標(biāo)查詢(xún)對(duì)應(yīng)一個(gè)掩模，標(biāo)簽和 ID 。根據(jù)不同的任務(wù)設(shè)置，和的格式和范圍是不同的。然而，和的格式和范圍是相似的。因此，將所有這些任務(wù)放入一個(gè)共享的編碼器和解碼器框架中，并共同訓(xùn)練一個(gè)模型用于所有分割任務(wù)是很自然的。

OMG-Seg架構(gòu)（具體是怎么做的～～）

OMG-Seg遵循Mask2Former的架構(gòu)設(shè)計(jì)。如圖所示，它包含一個(gè)骨干網(wǎng)絡(luò)、一個(gè)像素解碼器和一個(gè)掩模解碼器。區(qū)別在于以下幾個(gè)方面，包括凍結(jié)骨干的設(shè)計(jì)、合并的目標(biāo)查詢(xún)（包括目標(biāo)查詢(xún)和視覺(jué)提示），以及共享的多任務(wù)解碼器。根據(jù)不同的任務(wù)設(shè)置，解碼器輸出相應(yīng)的掩模和標(biāo)簽。

使用VLM編碼器作為凍結(jié)骨干 為了實(shí)現(xiàn)開(kāi)放詞匯的識(shí)別，在骨干部分，本文采用凍結(jié)的CLIP視覺(jué)模型作為特征提取器。本文使用ConvNeXt架構(gòu)來(lái)自O(shè)penCLIP。給定圖像/視頻輸入，VLM編碼器提取多尺度的凍結(jié)特征，以供進(jìn)一步處理。

像素解碼器作為特征適配器 像素解碼器與Mask2Former相同，包含多階段的可變形注意力層。它將凍結(jié)特征轉(zhuǎn)換為具有相同通道維度的融合特征，其中是特征的層索引，是最高分辨率的特征。

合并的目標(biāo)查詢(xún) 如上所分析，每個(gè)目標(biāo)查詢(xún)代表一種掩模輸出。然而，從功能的角度來(lái)看，圖像、視頻和交互模式代表不同的屬性。對(duì)于圖像，目標(biāo)查詢(xún)側(cè)重于目標(biāo)級(jí)別的定位和識(shí)別。對(duì)于視頻，目標(biāo)查詢(xún)可能涉及時(shí)間一致性，例如同一目標(biāo)跨越不同幀。對(duì)于交互式分割，目標(biāo)查詢(xún)被迫定位特定區(qū)域。對(duì)于圖像和視頻輸入，本文采用目標(biāo)查詢(xún)來(lái)表示圖像掩模或被跟蹤的tube掩模。因?yàn)閮烧叨夹枰Z(yǔ)義標(biāo)簽。本文將它們稱(chēng)為語(yǔ)義查詢(xún)，。對(duì)于交互模式，按照SAM 的做法，本文采用提示編碼器將各種視覺(jué)提示編碼成與目標(biāo)查詢(xún)相同形狀的查詢(xún)。本文將其稱(chēng)為位置查詢(xún)，。因此，本文可以共享變壓器解碼器的相同接口。

共享的多任務(wù)解碼器 其主要操作是交叉注意力，接收合并的目標(biāo)查詢(xún) 和 ) 和圖像/視頻特征，并輸出精煉的目標(biāo)查詢(xún)。最終的掩模是通過(guò)精煉查詢(xún)和高分辨率特征的點(diǎn)積獲得的。對(duì)于圖像語(yǔ)義級(jí)別的任務(wù)，本文采用了與Mask2Former相同的過(guò)程。具體而言，使用多尺度特征進(jìn)行掩碼交叉注意力 [18]。是查詢(xún)，而是鍵和值。然后，對(duì)精煉的查詢(xún)應(yīng)用多頭自注意力（MHSA）層。

對(duì)于視頻任務(wù)，本文采用相同的交叉注意力設(shè)計(jì)。唯一的區(qū)別是金字塔特征沿著時(shí)間維度與3D位置嵌入連接，這是先前的工作中的默認(rèn)設(shè)置。合并的視頻特征和精煉的查詢(xún)用于預(yù)測(cè)tube掩模。

對(duì)于交互式分割，本文執(zhí)行相同的交叉注意力設(shè)計(jì)。然而，本文跳過(guò)自注意力，以避免在MHSA層中的掩模查詢(xún)之間發(fā)生交互，因?yàn)榻换ナ椒指钪魂P(guān)心輸入視覺(jué)提示區(qū)域。在獲取精煉的目標(biāo)查詢(xún)后，它經(jīng)過(guò)一個(gè)預(yù)測(cè)FFN，通常包括一個(gè)帶有ReLU激活層和線(xiàn)性投影層的3層感知器。所有的查詢(xún)都受到掩模分類(lèi)損失和掩模預(yù)測(cè)損失的監(jiān)督。解碼過(guò)程以級(jí)聯(lián)的方式進(jìn)行，每個(gè)特征金字塔有三個(gè)階段。

Training and Inference（多數(shù)據(jù)集&多任務(wù)放一起怎么訓(xùn)～～）

聯(lián)合圖像視頻數(shù)據(jù)集聯(lián)合訓(xùn)練 與首先在圖像數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練不同，本文的目標(biāo)是僅聯(lián)合訓(xùn)練所有分割任務(wù)一次。所有訓(xùn)練目標(biāo)都是三種不同情況下的一個(gè)實(shí)體標(biāo)簽和掩模。實(shí)體可以是物體、雜項(xiàng)、類(lèi)別無(wú)關(guān)的掩模及其相應(yīng)的標(biāo)簽。注意，具有相同ID 的實(shí)例掩模形成tube掩模。在訓(xùn)練過(guò)程中，本文在預(yù)測(cè)的實(shí)體掩模和地面實(shí)體掩模之間應(yīng)用匈牙利匹配，以將目標(biāo)查詢(xún)分配給視頻/圖像實(shí)體，然后監(jiān)督它們的預(yù)測(cè)掩模和分類(lèi)。為了避免跨數(shù)據(jù)集的分類(lèi)沖突，分類(lèi)器被替換為CLIP文本嵌入。最終的損失函數(shù)表示為。這里，是掩模分類(lèi)的交叉熵（CE）損失，和分別是分割的掩模交叉熵（CE）損失和Dice損失。

通用推斷 對(duì)于圖像分割，本文遵循Mask2Former的相同推斷過(guò)程。例如，對(duì)于 PS，本文根據(jù)排序后的分?jǐn)?shù)合并物體和雜項(xiàng)。分?jǐn)?shù)由 CLIP 文本嵌入生成。對(duì)于視頻分割任務(wù)，對(duì)于 VIS 和 VPS，為了生成實(shí)例ID，本文遵循先前的工作，使用查詢(xún)匹配而不是引入額外的跟蹤組件。對(duì)于 VOS 任務(wù)，本文在第一幀和其余幀之間采用掩模匹配。對(duì)于交互式分割任務(wù)，本文遵循原始的 SAM，通過(guò)提供框和點(diǎn)提示來(lái)獲取二進(jìn)制掩模。對(duì)于開(kāi)放詞匯分割任務(wù)，由于本文有一個(gè)凍結(jié)的 CLIP 編碼器，本文將掩模池化分?jǐn)?shù)和學(xué)習(xí)分?jǐn)?shù)與開(kāi)放詞匯嵌入合并。

結(jié)合任務(wù)以獲取更多應(yīng)用 由于本文的模型可以執(zhí)行各種分割任務(wù)，結(jié)合交互、開(kāi)放詞匯和圖像/視頻分割任務(wù)可以導(dǎo)致一些新的應(yīng)用。例如，本文可以結(jié)合交互和視頻分割，實(shí)現(xiàn)靈活的基于提示驅(qū)動(dòng)的視頻目標(biāo)分割。或者本文可以將交互分割與開(kāi)放詞匯設(shè)置結(jié)合起來(lái)，從而實(shí)現(xiàn)開(kāi)放詞匯交互分割。

實(shí)驗(yàn)結(jié)果：

數(shù)據(jù)集和評(píng)估指標(biāo) 與常規(guī)設(shè)置不同，本文的目標(biāo)是盡可能地探索在多個(gè)數(shù)據(jù)集上進(jìn)行聯(lián)合訓(xùn)練。本文使用COCO panoptic、COCO-SAM、VIPSeg 和Youtube-VIS-2019（YT-VIS19）作為訓(xùn)練數(shù)據(jù)集。除了封閉集測(cè)試，本文還包括使用Youtube-VIS-2021、ADE-20k和DAVIS-2017數(shù)據(jù)集進(jìn)行開(kāi)放詞匯（OV）推理，它們的注釋在訓(xùn)練過(guò)程中未被使用。COCO-SAM 是通過(guò)使用真實(shí)邊界框，掩模中心點(diǎn)作為視覺(jué)提示來(lái)創(chuàng)建的。注釋是通過(guò)COCO全景掩模獲得的。此外，本文還在表3中包括了多數(shù)據(jù)集設(shè)置，以驗(yàn)證本文的OMG-Seg的多數(shù)據(jù)集聯(lián)合訓(xùn)練的有效性。除了表2之外，本文添加了更多數(shù)據(jù)集，包括ADE-20k 和YT-VIS21 用于聯(lián)合訓(xùn)練。本文為每個(gè)數(shù)據(jù)集使用相應(yīng)的評(píng)估指標(biāo)，包括 PQ、掩模 mAP、VPQ、tube mAP、J&F和 mIoU。

實(shí)驗(yàn)細(xì)節(jié) 本文在MMDetection中實(shí)現(xiàn)本文的模型和所有其他基線(xiàn)。本文使用32個(gè)A100 GPU的分布式訓(xùn)練框架。每個(gè)小批次在每個(gè)GPU上有一張圖像。對(duì)于數(shù)據(jù)增強(qiáng)，本文采用大規(guī)模的抖動(dòng)，與先前的工作一樣，以構(gòu)建強(qiáng)大的基線(xiàn)。對(duì)于每個(gè)表中的所有模型，本文采用相同的訓(xùn)練步驟。本文使用OpenCLIP來(lái)初始化骨干網(wǎng)絡(luò)，并用其相應(yīng)的文本嵌入替換學(xué)到的分類(lèi)器。對(duì)于圖像輸入，將它們視為偽視頻，通過(guò)將兩個(gè)圖像及其掩模連接成一個(gè)。本文采用不同的采樣率來(lái)平衡每個(gè)數(shù)據(jù)集的訓(xùn)練樣本。本文報(bào)告凍結(jié)和訓(xùn)練后的骨干網(wǎng)絡(luò)的結(jié)果供參考。

Main Results

雖然在COCO圖像分割基準(zhǔn)上表現(xiàn)略有下降，但在VIPSeg數(shù)據(jù)集上取得了接近最先進(jìn)水平的結(jié)果；開(kāi)放詞匯視頻分割領(lǐng)域效果也不錯(cuò)。

基于K-Net和 Mask2Former進(jìn)行了多數(shù)據(jù)集聯(lián)合訓(xùn)練。聯(lián)合共同訓(xùn)練通常可以提高大多數(shù)視頻分割數(shù)據(jù)集的性能，從而實(shí)現(xiàn)了大幅度的模型參數(shù)減少（從1326M減少到221M）。這種改進(jìn)在三個(gè)VPS和VIS數(shù)據(jù)集中是一致的，無(wú)論是否凍結(jié)了骨干網(wǎng)絡(luò)。

結(jié)論：

本文介紹了第一個(gè)用于圖像、視頻、開(kāi)放詞匯和交互式分割的聯(lián)合共訓(xùn)練框架。本文的解決方案OMG-Seg是一個(gè)新穎而簡(jiǎn)單的框架，它使用統(tǒng)一的查詢(xún)表示和共享的解碼器來(lái)處理多樣化的任務(wù)。首次有可能訓(xùn)練一個(gè)單一的分割模型，能夠在十個(gè)不同的任務(wù)上表現(xiàn)出與特定任務(wù)模型相媲美的性能。這種方法顯著減少了參數(shù)大小和在各種應(yīng)用程序的模型設(shè)計(jì)中需要專(zhuān)門(mén)工程的需求。預(yù)計(jì)本文高效且多功能的框架將成為多任務(wù)和多數(shù)據(jù)集分割的穩(wěn)健基線(xiàn)。