1. 程式人生 > >Deep Watershed Transform for Instance Segmentation

Deep Watershed Transform for Instance Segmentation

Min Bai & Raquel Urtasun
UfT

1. 傳統的分水嶺演算法簡介

影象處理中的分水嶺演算法常用來做影象區域分割(segmentation),基本的思路是計算一張energy map來表示影象,其中物體區域的energy和其他非物體區域的energy差別較大,從而形成包絡物體的分水嶺,稱包絡線為watershed line,稱物體區域為catchment basins(一般假設物體區域的energy低)。如下圖所示:
|center | 200x0
分水嶺演算法以一種類似注水的方法找到物體的區域,並通過抑制相鄰catchment basins的水交匯,達到分割物體的目的。

分水嶺演算法通常採用梯度作為energy map,這樣的話basin就對應比較平滑的區域。但是分水嶺演算法的一大侷限在於其很容易導致過分割。採用marker-controlled的方法,選擇marker又是一個很heuristic的問題。

2. Deep Watershed Transform

作者的基本思路是利用網路來學習分水嶺演算法裡的energy map,目的是希望學習得到的watershed line就表示instance的邊界。這樣直接應用分水嶺演算法就可以實現instance segmentation了。下圖表示傳統分水嶺演算法使用梯度大小得到的energy map和作者學習得到的energy map的對比:
|center | 400x0
可以看到,上圖中一共有3個instance,如果使用梯度大小作為energy map,容易受到噪聲的干擾,從而分割出7個instance,採用學習得到的energy map,能夠更好的表示一個instance。

直接學習instance之間的邊界是比較困難的,作者的做法分為兩步,首先訓練一個Direction Network學習instance內部點到邊界的方向,然後根據這個direction map再訓練一個Watershed Transform Network學習instance內部點的energy level,最後拿energy map應用分水嶺演算法。

整體的網路框架如下所示:
|center | 800x0

2.1. Direction Network

該網路的輸入為RGB影象以及語義分割的map(加入分割的結果作為輸入是為了讓網路只關注目標區域,忽略背景區域)。
目標是學習energy下降的方向(instance邊緣的能量低,instance中心的能量高,所以下降的方向指向instance的邊界)

對於每一個pixel,作者用二維單位向量表示該下降方向的真值,如下:

up,gt=Dgt(p)|Dgt(p)|
其中Dgt(p)表示畫素p到其對應instance邊界的最短距離。
能量梯度方向的示意圖如下所示:
|center | 300x0

作者方向的真值沒有選取角度的原因是因為角度的數值和方向不具有連續等價的性質,比如0度和360度方向一致但是數值差太大。採用能量梯度方向還有一個好處是:如果兩個instance之間存在相互遮擋,那麼loss就會對邊界處的能量梯度方向比較敏感,即如果一個pixel分錯了instance,那麼能量梯度方向肯定是和真值相反的,這樣loss會變得很大。這種loss能強迫網路學習pixel級別的精確定位。

Direction network輸出的是一個2-channel的direction map,網路部分如上圖所示,採用的是修改的vgg模型。這部分網路訓練的loss如下所示:

ldirection=pPobjwp||cos1<up,gt,up,pred>||2
其中<>表示求兩個向量內積, cos1得到的是兩個向量的夾角,wp是和instance面積成反比的係數,Pobj表示每一個instance。整體的loss是角度誤差的平方和。
2.2. Watershed Transform Network

通過前面的direction network得到每個畫素點的能量梯度方向之後,wtn再根據這張二維的map學習每一個pixel對應的energy level。其中level 0表示背景以及和距離邊界2個pixel以內的畫素點。level越高,畫素點越靠近物體的中心。作者一共手動選擇了K = 16個level。

WTN模組的訓練過程如上所示,輸入為direction map,輸出為energy map。其訓練的loss為改動的交叉熵:

lwatershed=pPobjk=1Kwpck(t¯p,klogy¯p,k+tp,klogyp,k)
其中wp還是和原來一樣的和instance面積成負相關的係數,ck表示和energy level相關的係數,因為我們最後是選擇level大於0或1來分割instance,因此我們希望level小的部分學習儘量準確,所以ck對於小的k,值比較大。

整體訓練過程中各個步驟的輸入輸出如下所示:
|center | 800x0

2.3. Energy Cut and Instance Extraction

最後分割的時候,對於小目標類別(person, rider, etc),選擇level=0進行分割,對於大目標類別(car, truck, etc),採用level=1進行分割。

3. Experiments

作者只在CityScapes上做了實驗。
訓練的時候,首先先拿RGB影象,分割的gt以及instance segmentation的gt訓練direction network;
然後拿gt的direction map和gt的instance segmentation訓練WTN。
最後把兩個模組級聯在一起,用RGB影象和PSPNet的輸出做fine-tune。

在CityScapes上比之前的instance segmentation效果好很多
|center | 700x0

網路對於分割的效果也不是特別敏感,但是分割好,instance segmentation也會更好
|center | 400x0

定性分析結果如下:
|center | 900x0