[InstanceSegmentation] SOLO:Segmenting Objects by Locations 논문 정리

오늘 정리할 논문은 one stage instance segmentation 네트워크 인 SOLO이다. 

Paper:

SOLO: Segmenting Objects by Locations

SOLO: Segmenting Objects by Locations

이 논문은 e2e로 학습도 가능 하고 instance를 구분 하는 방식이 매우 흥미 롭다.

 

SOLO의 접근 방식:

SOLO에서 해결하고자 하는 건 instance segmentation 이다. 

이미지 상에 존재하는 여러 물체를 각각 구분해 mask를 생성하는 문제이다. 

 

이를 위해 SOLO에서는 object center와 size정보를 이용한다. 

bbox 예측이나 post processing을 통한 instance pixel grouping없이 개별 object를 위 언급한 object center와 size를 이용해 정의한 center region와 해당 region에 대응 하는 mask로 직접 instance mask를 생성한다. 

 

간단히 기존 instance segmentation 네트워크의 전략을 살펴 보면 크게 top-down, bottom-up 방식으로 나눌 수 있다.

• top-down 방식:  bbox 찾고 해당 bbox 안에 있는 object mask를 생성. 이 방식에서 instance의 구분은 bbox거의 전적으로 의존하게 된다. 
  대표적인 모델은 Mask-RCNN, PANet(mask-rcnn에서 feature representation을 강화한 방식, Mask-scroing rcnn(mask prediction의 품질을 평가하기 위한 mask-iou branch를 도입)
• bottom up 방식: 이 방식의 네트워크 들은 대부분 각각의 object를 구분하기 위해 bbox를 사용하는 대신
  각 object를 구분할 수 있는 임의의 latent space에서 각 object를 나타내는 pixel embedding(의 형태는 다양 할 수 있다.)을 학습하고 post processing을 통해 grouping 한다. 
  대표적인 모델은 
  1. SGN: instance segmentation을 단계적인 sub-grouping문제로 정의. break point 라는 개념을 도입해 각 instance의 경계에 있는 point를 찾고 이 point를 line segment로 line segment를 instance를 grouping해 instance를 구성한다. 
  2. SSAP(single shot instance segmentation with affinity pyramid): pixel이 같은 인스턴스에 속하는 지를 pixel pair affinity라는 확률 값을 계산 하는 문제로 정의

SOLO는 bbox를 찾지도 pixel embedding을 학습해 후처리로 그룹핑 하지도 않고 mask 를 바로 생성하기 떄문에 

Direct mask prediction 이라고 할 수 있다. 

 

SOLO Architecture:

Fig 1. Solo framework. grid size=5 일 때

Fig 1. 은 매우 직관적으로 solo의 framework을 이해 할 수 있다. 

전체 흐름을 간단히 보면 입력 이미지를 SxS grid로 나타내고 어떤 object의 center가 특정 grid cell에 위치하면

그 cell이 sementic category와 instance mask를 예측하는 책임을 진다. 

 

Sementic Category:

sementic category는 object의 class를 예측하는 branch이다.

Object class 가 c개 있다면 category branch는 $ S \times S \times C$ 형태의 tensor를 생성한다.  이때 grid cell (i,j)에 있는 길이 c의 vector는 해당 cell에 object center가 존재하는 객채가 어떤 class에 속할지 에 대한 확률값 벡터이다. 

 

Instance mask:

Instance mask head의 output shape은 $ H \times W \times S^2 $  이다.

채널이 $ S^2 $ 인 이유는 grid cell (i,j) 각각에 대응 하는 채널이 하나씩 있기 때문이다. 

즉 Intance mask의 각 채널은 category branch의 grid cell (i,j)와 1:1 매핑 되는 관계이다.

예를 들어 임의의 객체 A의 center가 grid cell (i,j)에 위치 한다면 이 객체의 segmentation mask는 

$ mask channel = i \times S + j$ intance mask 에 생성되게 끔 네트워크를 학습한다. 

 

Fig 1.의 input image와 category branch 를 보면 왼쪽 얼룩말의 센터는 grid cell에서 붉은 색$(i_{1}, j_{1})$으로, 오른쪽 얼룩 말의 센터는 grid cell에서 파란색$(i_{2}, j_{2})$으로 표시 된다. 그리고 이 두 얼룩말의 마스크는 각 grid cell에 대응 하는 intance mask channel $ i_{1} \times S + j_{1}$, $ i_{2} \times S + j_{2}$에 각각 생성되는 식이다. 

 

이렇게 grid cell 위치와 instance mask 채널이 서로 연관되어 있기 때문에 SOLO는 "coordinate aware" convolution을 사용 했다고 한다.(기존 convolution은 translation invarian한 특성을 가지고 있기 떄문에 이와는 구분된다.)

그리고 이를 위해 CoordConv라는 convolution을 사용 했다. 

CoordConv는 instance mask branch의 입력과 같은 spatial size를 갖는 $ H \times W \tmes 2$ 크기의 텐서에 coordinate을  [-1,1]로 normalize하고 backbone의 output에  concat해서 instance mask branch에서 사용 한다. 

 

Fig 2의 Mask branch에서 input channel이 (256+2)인 이유가 바로 이 coordinate 정보를 concat했기 때문 이다. 

 

아래 코드를 보면 쉽게 이해가 된다. 

$코드 입력$

 

Fig 2. SOLO category and Mask branch

FPN:

SOLO는 feature pyramid layer(FPN)의 output이 각 category branch와 mask branch의 입력으로 사용되어 서로 다른 크기의 object를 찾는 효율을 높였다. 

 

학습:

GT 생성: 

 intance mask를 category classification grid cell위치에 대응 하는 instance mask channel에 binary mask 형태로 생성한다. 따라서 학습시 instance mask를 이에 대응 하는 형태로 만들어야 한다.

 

 앞서 object center가 (i, j) grid cell에 위치한다고 할때 해당 cell 이 object의 class와 mask생성을 담당한다고 했다. 이때, object size에 따라 여러 cell이 한 object의 class와 mask생성을 담당할 수도 있다. object center 좌표는 object mask gt 의 centre of mass로 정의하고 object 크기를 고려해 center sampling을 하기 위해 center region은 $ (C_{x}, C_{y}, \epsilon W, \epsilon H) \epsilon=0.2$ 로 정의한다. 

 

 이렇게 center region을 정하고 나면 category grid cell 상에 아래와 같이 center region이 형성될 수 있는데 이때 각 cell 위치에 해당하는 instance mask channel의 binary mask가 똑같이 들어간다. 

 

Fig 3은 이해를 돕기 위해 그린 그림이다. 붉은 색영역은 임의의 object에 대한 center region 이고 이에 대응 하는 mask gt가 intance mask gt에 어떻게 저장되는지 보여 준다. 

Fig 3. center region 과 instance mask gt 예시

LOSS:

$$ total loss = L_{category} + \lambdaL_{masks} $$

Category classification($L_{category}): 

solo는 category classification을 위해 focal loss를 사용 했다. (굳이 식을 설명하지 않아도 될정도로 많이 사용되니 식은 생략한다.)

 

Segmentation Mask($L_{masks}$):

Dice loss:

$ L_{mask}= \frac{1}{N_{pos}} \sum I(p_{i,j} >0) d_{mask}(m_{k},m_{k}^{*}) $

$ d_{mask}=1-\frac{2\sum_{x,y}(p_{x,y},q_{x,y}}{\sum_{x,y}p_{x,y}^{2}+\sum_{x,y}q_{x,y}^{2}} $

 

$p_{i.j}$ 는 category target으로 i,j는 grid cell 상의 좌표.

$m_{k}^{*}$는 mask gt

$d_{mask}$에서 $p_{x,y}$는 mask prediction 에서 x,y위치의 pixel value, $q_{x,y}$는 gt mask의 pixel value이다. 

 

 

Inference:

inference 시 category score threshold는 0.1을 사용해 낮은 score를 필터링 하고 score상위 500개에 대응되는 mask에 NMS를 적용한다. 

최종 score를 계산하기 위해선 maskness를 category score에 곱해 주는데 maskness란 $\frac{1}{N_{f}\sum p_{i}}$로 표현되고 $p_{i}$는 prediction mask의 pixel value 다. 

(mask ness 를 곱해주는 의미는 maskness점수가 높다는게 곧 해당 object의 형태를 잘 찾았다는 것일 태니 그걸 최종 score에 반영하는 것 같다. )

 

결과:

Fig 4. COCO 데이터셋에 대한 solo 학습 결과

Fig 4.에서 볼 수 있듯이 SOLO는 Mask R-CNN과 비슷 하거나 더 좋은 결과를 보여 준다. one stage instance segmenation이 two stage보다 좋아 질 수 있다는 충분한 가능성을 보여주는 예시 인거 같다. 

가장 좋은 성능을 보인것은 Resnet 101 backbone+ deformable convolution+ FPN을 사용 한 모델이다. 

 

 

Fig 5. SOLO behavior 설명

Fig 5에서는 SOLO가 어떻게 동작 하는지를 볼 수 있는데 아랫 줄의 보라색은 category branch 의 grid 가 12 x12일때

대응 되는 instance mask branch의 각 채널이 어떤 prediction결과를 가지는 지 보여준다. 

 instance mask branch output의 각 채널은 하나의 object에 대응 하는 mask를 생성하는 것을 볼 수 있는데 이 mask는 해당 채널과 대응 하는 category branch의 grid cell에 center가 위치하는 객채의 마스크 이다. 

 

객체의 크기에 따라 center region 이 여러 cell을 포함 할 수 있기때문에 대응 되는 intance mask channel도 여러개 일수 있어서 한  object의 mask가 여러개 생성되는 모습을 볼 수 있는데 NMS를 이용해 같은 객체의 mask를 하나로 합친다. 

 

 

Fig 6. FPN grid size의 영향성

Fig 6은 category branch의 grid size에 따른 AP 결과를 보여 준다. grid number 12,24,36은 단일 grid size를 사용 했을 경우를 말하고 Pyramid는 FPN을 사용해 여러 grid size를 사용해 결과를 추정 했을때를 말한다. 

단일 size이 grid를 사용할 경우 grid size가 클수록 AP가 높게 나왔고 FPN을 사용해 여러 grid size를 고려 하는게 가장 높은 결과를 모여준다. FPN의 경우 다양한 object 크기를 고려 할 수 있으니 당연한 결과로 생각된다. 

Fig 7. CoordConv 영향성

Fig 7은 coordinate aware convolution 이 결과에 어떤 영향을 미치는지에 대한 테스트 이다. 

Coordconv 를 사용 하지 않을 경우 Fig 7 테이블의 가장 윗줄의 AP =32.2가 coorconv를 사용했을 때보다 최소 3.6point가 낮은 결과를 보여 준다. 비슷 하게 coordconv를 2회 이상 사용하는 것은 성능에 영향을 거의 미치지 않는 것을 볼 수 있기 때문에 논문에서는 2회만 사용 했다. 

 

논문에서는 좀더 많은 ablation experiment를 진행했지만 이정도만 정리해도 충분할거 같다. 

 

coordConv의 내용이 좀 흥미로워서 이 논문을 한번 읽어봐야겠다.