牛津中國(guó)小哥提出“3D-BoNet”,比3D點(diǎn)云實(shí)例分割算法快10倍!
本文提出了一種基于邊界框回歸的高效點(diǎn)云實(shí)例分割算法,通過(guò)最小化關(guān)聯(lián)代價(jià)函數(shù)來(lái)實(shí)現(xiàn)大致的邊界框回歸,并通過(guò)point mask預(yù)測(cè)來(lái)實(shí)現(xiàn)最終的實(shí)例分割。3D-BoNet不僅在ScanNet和S3DIS數(shù)據(jù)集上達(dá)到了state-of-the-art的效果,比當(dāng)前大多數(shù)算法快了10倍以上。
Introduction
實(shí)現(xiàn)有效的三維場(chǎng)景理解(3D scene understanding)是計(jì)算機(jī)視覺(jué)和人工智能領(lǐng)域的關(guān)鍵問(wèn)題之一。近年來(lái),針對(duì)三維點(diǎn)云理解的研究取得了顯著的進(jìn)展,在諸如點(diǎn)云目標(biāo)檢測(cè),語(yǔ)義分割等任務(wù)上都展現(xiàn)出了很不錯(cuò)的效果。然而,針對(duì)于點(diǎn)云實(shí)例分割的研究還處于較為初級(jí)的階段。
Motivation
如下圖所示,當(dāng)前主流的點(diǎn)云實(shí)例分割算法可以分為以下兩類(lèi):1)基于候選目標(biāo)框(Proposal-based methods)的算法,例如3D-SIS[1],GSPN[2],這類(lèi)方法通常依賴(lài)于兩階段的訓(xùn)練(two-stage training)和昂貴的非極大值抑制(non-maximum suppression, NMS)等操作來(lái)對(duì)密集的proposal進(jìn)行選擇。2)無(wú)候選目標(biāo)框的算法(Proposal-free methods),例如SGPN[3], ASIS[4], JSIS3D[5], MASC[6], 3D-BEVIS[7]等。這類(lèi)算法的核心思想是為每個(gè)點(diǎn)學(xué)習(xí)一個(gè)discriminative feature embedding,然后再通過(guò)諸如mean-shift等聚類(lèi)(clustering)方法來(lái)將同一個(gè)instance的點(diǎn)聚集(group)到一起。這類(lèi)方法的問(wèn)題在于最終聚類(lèi)到一起的instance目標(biāo)性(objectness)比較差。此外,此類(lèi)方法后處理步驟(post-processing)的時(shí)間成本通常較高。
圖1. 當(dāng)前主流的點(diǎn)云實(shí)例分割算法對(duì)比
不同于上述兩類(lèi)方法,我們提出了一個(gè)single stage, anchor free并且end-to-end的基于邊界框回歸的實(shí)例分割算法(3D-BoNet)。該算法具有如下優(yōu)勢(shì)
- 相比于proposal-free的方法,3D-BoNet顯式地去預(yù)測(cè)目標(biāo)的邊界框,因此最終學(xué)到的instance具有更好的目標(biāo)性(high objectness).
- 相比于proposal-based的方法,3D-BoNet并不需要復(fù)雜耗時(shí)的region proposal network以及ROIAlign等操作,因此也不需要NMS等post-processing步驟。
- 3D-BoNet由非常高效的shared MLP組成,并且不需要諸如非極大值抑制,特征采樣(feature sampling),聚類(lèi)(clustering)或者投票(voting)等后處理步驟,因此非常高效。
Overview
3D-BoNet的總體框架如下圖所示,它主要由Instance-level bounding box prediction和Point-level mask prediction兩個(gè)分支組成。顧名思義,bounding box prediction分支用于預(yù)測(cè)點(diǎn)云中每個(gè)實(shí)例的邊界框,mask prediction分支用于為邊界框內(nèi)的點(diǎn)預(yù)測(cè)一個(gè)mask,進(jìn)一步區(qū)分邊界框內(nèi)的點(diǎn)是屬于instance還是背景。
圖2. 3D-BoNet的總體框架
看到這里,你可能會(huì)產(chǎn)生疑惑:這個(gè)看起來(lái)跟proposal-based的框架好像也沒(méi)什么區(qū)別?
先說(shuō)結(jié)論:區(qū)別很大。但問(wèn)題是區(qū)別到底在哪里呢?
首先,我們可以回顧下proposal-based方法是怎么產(chǎn)生邊界框的呢?沒(méi)錯(cuò),就是根據(jù)anchor用region proposal network (RPN)來(lái)產(chǎn)生大量密集的邊界框然后再進(jìn)一步refine,但這個(gè)顯然不夠高效,而且是否真的有必要產(chǎn)生這么多密集的邊界框呢?針對(duì)這個(gè)問(wèn)題,我們可以來(lái)一個(gè)大膽的假設(shè):要不不用RPN,直接讓為每一個(gè)instance回歸(regress)一個(gè)唯一的,但可能不是那么準(zhǔn)確的邊界框呢(如圖3所示)?
圖3. 為每一個(gè)instance回歸一個(gè)大致的邊界框示例
考慮到3D點(diǎn)云本身就顯式地包含了每個(gè)物體的幾何信息,我們認(rèn)為這個(gè)目標(biāo)是可行的。然后再更大膽一點(diǎn),要不直接用global feature來(lái)regress每個(gè)instance的邊界框試試?如果能做到這點(diǎn),那問(wèn)題不就解決一半了嗎?
但新的問(wèn)題馬上又來(lái)了。。首先,每個(gè)三維場(chǎng)景中所包含的實(shí)例數(shù)目是不一樣的(如何讓網(wǎng)絡(luò)自適應(yīng)的輸出不同個(gè)數(shù)的邊界框?),而且每個(gè)點(diǎn)云中的實(shí)例還是無(wú)順序的。這就意味著我們即便用網(wǎng)絡(luò)regress了一系列邊界框,也難以將這些邊界框和ground truth的邊界框一一對(duì)應(yīng)的聯(lián)系起來(lái),進(jìn)一步帶來(lái)的問(wèn)題就是:我們無(wú)法實(shí)現(xiàn)對(duì)網(wǎng)絡(luò)的有監(jiān)督的訓(xùn)練和優(yōu)化。
到這里,核心的問(wèn)題就變成了:我們應(yīng)該怎么去訓(xùn)練這種網(wǎng)絡(luò)呢?
針對(duì)這個(gè)問(wèn)題,我們提出了邊界框關(guān)聯(lián)層(bounding box association layer)以及multi-criteria loss 函數(shù)來(lái)實(shí)現(xiàn)對(duì)網(wǎng)絡(luò)的訓(xùn)練。換句話說(shuō),我們要把這個(gè)預(yù)測(cè)邊界框和ground truth邊界框關(guān)聯(lián)(配對(duì))的問(wèn)題建模為一個(gè)最優(yōu)分配問(wèn)題。
圖4. 邊界框預(yù)測(cè)分支的結(jié)構(gòu)圖
如何關(guān)聯(lián)?
為了將網(wǎng)絡(luò)預(yù)測(cè)出來(lái)的每一個(gè)邊界框與ground truth 中的邊界框唯一對(duì)應(yīng)地關(guān)聯(lián)起來(lái),我們將其建模為一個(gè)最優(yōu)分配問(wèn)題。假定 是一個(gè)二值(binary)關(guān)聯(lián)索引矩陣,當(dāng)且僅當(dāng)?shù)?個(gè)預(yù)測(cè)的邊界框分配給ground truth的邊界框時(shí)。 是關(guān)聯(lián)代價(jià)矩陣, 代表將第 個(gè)預(yù)測(cè)的邊界框分配給ground truth的邊界框的關(guān)聯(lián)代價(jià)。一般來(lái)說(shuō),代表兩個(gè)邊界框的匹配程度,兩個(gè)邊界框越匹配也即代價(jià)越小。因此,邊界框的最優(yōu)關(guān)聯(lián)問(wèn)題也就轉(zhuǎn)變成為尋找總代價(jià)最小的最優(yōu)分配索引矩陣 的問(wèn)題,用公式表示如下:
下一步,如何計(jì)算關(guān)聯(lián)代價(jià)矩陣?
衡量?jī)蓚€(gè)三維邊界框的匹配程度,最簡(jiǎn)單直觀的評(píng)價(jià)指標(biāo)就是比較兩個(gè)邊界框的最小-最大頂點(diǎn)之間的歐式距離。然而,考慮到點(diǎn)云通常都非常稀疏且不均勻地分布在3D空間中,如圖4所示,盡管候選框#2(紅色)與候選框#1(黑色)與ground truth邊界框#0(藍(lán)色)都具有相同的歐式距離,但框#2顯然具有更多有效點(diǎn)(overlap更多)。因此,在計(jì)算代價(jià)矩陣時(shí),有效點(diǎn)的覆蓋率也應(yīng)該要考慮進(jìn)來(lái)。
圖5. 預(yù)測(cè)邊界框與真實(shí)邊界框點(diǎn)云覆蓋率示意圖
為此,我們考慮以下三個(gè)方面的指標(biāo):
(1) 頂點(diǎn)之間的歐式距離。舉例來(lái)說(shuō),第個(gè)預(yù)測(cè)的邊界框 分配給ground truth的邊界框 的代價(jià)為:
(2) Soft IoU。給定輸入點(diǎn)云 以及第 ground truth的實(shí)例邊界框 ,我們可以直接得到一個(gè)hard的二值(binary)矢量 來(lái)表征每個(gè)點(diǎn)是否在邊界框內(nèi)。然而,對(duì)于相同輸入點(diǎn)云的第 預(yù)測(cè)框,直接獲得類(lèi)似的hard的二值矢量將導(dǎo)致框架不可微(non-differentiable)。因此,我們引入了一個(gè)可微但簡(jiǎn)單的算法來(lái)獲得類(lèi)似但soft的二值矢量,稱(chēng)為point-in-pred-box-probability,詳情見(jiàn)paper Algorithm 1。 所有值都在 范圍內(nèi),這個(gè)值越高代表點(diǎn)在框內(nèi)可能性越大,值越小則對(duì)應(yīng)點(diǎn)可能離框越遠(yuǎn)。因此,我們定義第個(gè)預(yù)測(cè)的邊界框以及第ground truth的邊界框的sIoU如下:
(3) 此外,我們還考慮了和
之間的交叉熵。交叉熵傾向于得到更大且具有更高覆蓋率的邊界框:
總結(jié)一下,指標(biāo)(1)盡可能地使學(xué)到的框與ground truth的邊界框盡可能地重合,(2)(3)用于盡可能地覆蓋更多的點(diǎn)并克服如圖5所示的不均勻性。第 個(gè)預(yù)測(cè)的邊界框與第ground truth的邊界框的最終關(guān)聯(lián)代價(jià)為:
Loss function如何定義?
在通過(guò)邊界框關(guān)聯(lián)層之后,我們使用關(guān)聯(lián)索引矩陣 對(duì)預(yù)測(cè)邊界框 及其對(duì)應(yīng)分?jǐn)?shù) 與groundtruth進(jìn)行匹配,使得靠前的 個(gè)邊界框(ground truth的總邊界框數(shù))及與ground truth的邊界框能匹配上。
針對(duì)邊界框預(yù)測(cè)我們采用了多準(zhǔn)則損失函數(shù),也即三者求和:
針對(duì)邊界框分?jǐn)?shù)預(yù)測(cè)我們采用了另外一個(gè)損失函數(shù)。預(yù)測(cè)框分?jǐn)?shù)旨在表征相應(yīng)預(yù)測(cè)框的有效性。在通過(guò)關(guān)聯(lián)索引矩陣 重新排序以后,我們?cè)O(shè)定前 個(gè)真實(shí)的邊界框?qū)?yīng)的分?jǐn)?shù)為1,剩余的無(wú)效的 個(gè)邊界框?qū)?yīng)的分?jǐn)?shù)為 0。我們對(duì)這個(gè)二元分類(lèi)任務(wù)使用交叉熵?fù)p失
作為另外一個(gè)并行的分支,我們的方法可以采用任意現(xiàn)有的點(diǎn)云語(yǔ)義分割算法(比如Sparseconv, Pointnet++等等)作為對(duì)應(yīng)的語(yǔ)義分割模塊,整個(gè)網(wǎng)絡(luò)最終的loss function定義為
代表語(yǔ)義分割分支的loss,這里我們采用標(biāo)準(zhǔn)的交叉熵。網(wǎng)絡(luò)具體的優(yōu)化和求解過(guò)程我們采用Hungarian算法,詳情請(qǐng)見(jiàn)[8],[9]。
如何預(yù)測(cè)instance mask?
相比于bounding box prediction分支,這個(gè)分支就相對(duì)簡(jiǎn)單很多了,因?yàn)橹灰吔缈蝾A(yù)測(cè)的足夠好,這個(gè)分支就相當(dāng)于做一個(gè)二值分類(lèi)問(wèn)題,即便瞎猜也能有50%正確率。在這個(gè)分支中,我們將點(diǎn)的特征點(diǎn)與每個(gè)邊界框和分?jǐn)?shù)融合在一起,隨后為每一個(gè)實(shí)例預(yù)測(cè)一個(gè)點(diǎn)級(jí)別的二值mask??紤]到背景點(diǎn)與實(shí)例點(diǎn)數(shù)的不平衡,我們采用focal loss[10] 對(duì)這個(gè)分支進(jìn)行優(yōu)化。
圖 6. Point mask prediction分支結(jié)構(gòu)圖。
Experiments
在ScanNet(v2) benchmark上,我們的方法達(dá)到了state-of-the-art的效果,相比于3D-SIS,MASC等方法都有顯著的提升。
圖7. 我們的方法在ScanNet(V2)的結(jié)果
在Ablation study中,我們也進(jìn)一步證實(shí)了各個(gè)分支以及l(fā)oss function各個(gè)評(píng)估指標(biāo)的作用。詳細(xì)的分析見(jiàn)paper。
圖 8. Ablation study結(jié)果 (S3DIS, Area5)
在計(jì)算效率方面,3D-BoNet是目前速度最快的方法,相比于SGPN, ASIS, 3D-SIS等方法,3D-BoNet快了十倍以上。
圖 9. 不同方法處理ScanNet validation set所需要的時(shí)間消耗。
此外,我們還在圖10中進(jìn)一步展示了我們提出的loss function在S3DIS數(shù)據(jù)集進(jìn)行訓(xùn)練時(shí)(Area1,2,3,4,6進(jìn)行訓(xùn)練,Area 5進(jìn)行測(cè)試)的變化曲線。從圖中可以看到,我們提出的loss function能夠比較一致的收斂,從而實(shí)現(xiàn)對(duì)語(yǔ)義分割分支,邊界框預(yù)測(cè)分支以及point mask預(yù)測(cè)分支端到端方式的優(yōu)化。
圖10. 我們的方法在S3DIS數(shù)據(jù)集上的training loss
在圖11中我們給出了預(yù)測(cè)邊界框和邊界框分?jǐn)?shù)的可視化結(jié)果。可以看出,我們方法預(yù)測(cè)出來(lái)的框并不一定非常精準(zhǔn)和緊湊。相反,它們相對(duì)比較松弛(inclusive)并且具有比較高的目標(biāo)性(high objectness)。這也與本文一開(kāi)始希望得到的大致邊界框的目標(biāo)相一致。
圖11. 我們的方法在S3DIS數(shù)據(jù)集Area 2上的預(yù)測(cè)邊界框和分?jǐn)?shù)的可視化。紅色框表示預(yù)測(cè)的邊界框,藍(lán)色的邊界框代表ground truth。
當(dāng)邊界框已經(jīng)預(yù)測(cè)好以后,預(yù)測(cè)每個(gè)框內(nèi)的point mask就容易很多了。最后我們可視化一下預(yù)測(cè)的instance mask,其中黑點(diǎn)代表屬于這個(gè)instance的概率接近為0,而帶顏色的點(diǎn)代表屬于這個(gè)instance的概率接近為1,顏色越深,概率越大。
圖12. 預(yù)測(cè)instance mask的可視化結(jié)果。輸入點(diǎn)云總共包含四個(gè)instance,也即兩個(gè)椅子,一個(gè)桌子以及地面。從左到右分別是椅子#1,椅子#2,桌子#1,地面#2的point mask.
最后總結(jié)一下,我們提出了一種基于邊界框回歸的高效點(diǎn)云實(shí)例分割算法,通過(guò)最小化匹配代價(jià)函數(shù)來(lái)實(shí)現(xiàn)大致的邊界框回歸,并通過(guò)point mask預(yù)測(cè)來(lái)實(shí)現(xiàn)最終的實(shí)例分割。我們提出的3D-BoNet不僅在ScanNet和S3DIS數(shù)據(jù)集上達(dá)到了state-of-the-art的效果,并且比現(xiàn)有其他算法更加高效。