[Vision] Improving Pixel Embedding Learning through Intermediate Distance Regression Supervision for Instance Segmentation

논문 정보

제목 : Improving Pixel Embedding Learning through Intermediate Distance Regression Supervision for Instance Segmentation

저자 : Yuli Wu 외

시기 : 13 July 2020

 

논문 세 줄 요약

Distance Regression Supervision의 도움으로 Embedding vector를 잘 학습해서 Instance Segmentation을 잘 해보자. ECCV의 Leaf Segmentation 워크샵 (The CVPPP 2020)에서 Best dice 기준 1등을 했고 Best Paper Honorable Mention 상을 탔습니다.

 

기초 지식

1. Embedding

Embedding은 자연어 처리 (NLP)에서 많이 보셨을텐데요 자연어 처리에서 임베딩은 단어 하나하나를 벡터로 표현하기 위해 사용됩니다. 단어를 n개의 글자가 조합된 문자열로 보는 것보다는 단어 각각이 k차원의 벡터를 가질 수 있도록 표현할 수 있다면 (이걸 뉴럴 넷으로 하죠) 더 쉽게 단어들, 문장들을 다룰 수 있습니다. 예를 들어 학습을 통해서 비슷한 뉘앙스의 단어는 임베딩 벡터가 비슷하게 표현할 수 있다면 단어들 사이의 거리도 계산할 수 있습니다.

이 논문에서는 Pixel Embedding이라는 워딩을 사용합니다. 사실 정말 특별한 것이 아닙니다. 만약 RGB 이미지라면 x, y 픽셀 하나하나를 3차원 벡터로 볼 수 있을텐데요, 어떠한 방식으로 (CNN을 써서) (x, y, 3)인 이미지가 (x, y, N)의 데이터가 된다면 픽셀 하나하나가 N차원의 임베딩 벡터를 가지게 되는 것이죠. 때문에 Pixel Embedding은 특별한 내용이 아니라 각 픽셀 당 N차원 벡터를 가질 수 있도록 CNN을 설계하고 이 임베딩 벡터를 가지고 잘 다뤄서 학습을 해보자는 것입니다. 

 

2. Cosine Similarity

Cosine similarity는 정말 단순한데, 그저 두 벡터 사이의 각도의 코사인 값입니다. 일반적인 경우에 이 cosine similarity는 양수 공간에서만 사용합니다. 양수 공간에서 사용하게 되면 이 similarity를 작게 만드는 것이 180도가 아니라 90도가 되기 때문에 벡터들 간의 orthogonality를 강제할 수 있기 때문입니다.

 

메인 논문 내용

이것이 이들이 제안한 아키텍쳐입니다.

총 2개의 U-Net이 있고, 2개의 loss 함수가 있습니다. 

U-Net은 input과 output의 크기가 같은 네트워크인데 풀링과 업샘플링을 통해서 굉장히 큰 Field Of View (FOV)를 가져갈 수 있습니다. 이곳에서는 'Pixel' embedding이기 때문에 x, y를 유지시키는 네트워크인 U-Net을 사용한 것 같습니다. 물론 U-Net이 정말 널리 사용되고 있기도 하고, FOV도 크고 여러 장점들이 있어서 선택한 이유도 있을 것이고요. 

그렇다면 Loss 함수를 보자면,

1. D-Loss

이는 Distance loss의 약자인데, Distance regression을 합니다. 자세히 설명하자면, 학습 데이터에서 각 instance의 정보가 있으면 OpenCV의 distance transform을 통해서 각 픽셀 당 가장 인접한 경계까지의 거리를 픽셀값으로 만듭니다. 예시를 보여드리자면,

말의 모양이 있고 Distance Transform을 하면, 위와 같이 픽셀 하나하나당 가까운 경계까지의 거리를 그릴 수 있습니다. 위에 있는 U-Net 1은 Leaf 이미지를 받아와 채널 1짜리 Distance Map을 맞춰야 합니다. 그래서 Distance Map과의 MSE loss로 Regression task를 수행합니다. 논문의 제목에서는 "through intermediate distance regression supervision" 라는 워딩으로 되어 있는데요, 여기서 알 수 있다시피 Distance Map으로 학습하는 U-Net 1은 메인이 아닙니다. 이 모델이 학습한 것은 이제 말할 U-Net 2에 도움을 주는 애입니다.

 

2. E-Loss

이곳이 Pixel Embedding이 사용되는 곳입니다. 다른 U-Net 2은 output으로 (x, y, E)를 내보냅니다. 여기서 E가 앞서 말한 임베딩 벡터의 차원이고 이 논문에서는 4, 8, 16, 32 등 여러 값으로 테스트해보았습니다. 

E-Loss는 Cosine similarity를 통해 구해집니다. 위에서 설명한대로 cosine similarity는 결국 두 벡터 사이의 각도입니다. E-loss는 두 개의 cosine similarity term이 있는데요, inter loss와 intra loss입니다.

뮤_cA는 인스턴스 cA의 각 픽셀별로 임베딩 벡터의 평균입니다. e_i는 픽셀 i 하나의 임베딩 벡터입니다.

복잡하게 생겼지만 쉽게 생각하면 inter loss는 두 개의 인스턴스가 각도가 커지도록 만드는 loss입니다. cosine 함수를 생각해보면 값이 작아지려면 0도보다는 90도가 좋겠죠. Intra loss는 인스턴스 내부에서 각 픽셀들이 인스턴스 픽셀들의 평균 벡터와 같은 방향을 가리키도록 만드는 loss 입니다.

쉽게 생각하면 다른 인스턴스끼리는 다른 방향, 같은 인스턴스 내에서는 픽셀들이 같은 방향을 띄도록 만들어줍니다.

 

이렇게 두 개의 loss, 임베딩 벡터를 만드는 loss와 distance map을 regression을 통해 학습하는 loss가 있습니다.

 

3. Post processing

이후 이제 instance를 후처리를 통해 각각 구분해내야 하는데,

1. distance map에서 max값의 몇 퍼센트로 thresholding을 해서 몇 개의 seed point를 잡는다. distance map은 결국 인스턴스의 중심으로 갈 수록 값이 커지니까 인스턴스의 중심들을 seed point들로 만듭니다.

2. seep point를 기준으로 embedding vector를 봐서 임베딩 벡터가 45도 이내에 속하는 것들로 모아 하나의 instance로 만듭니다. 여러 각도로 테스트해봤는데 45도가 제일 좋았다고 합니다. 이 angular clustering은 앞서 설명드린 E-loss가 같은 instance면 같은 방향, 다른 instance면 다른 방향을 띄도록 학습을 했기 때문에 가능하겠죠.

 

 

특이사항

보면 위의 도식에서 D-feat가 이미지와 concat 된 후에 U-Net 2에 들어가는 것을 알 수 있습니다. 이 논문에서 말하길 임베딩 벡터를 그냥 학습하는 것은 잘 안되는데 distance map을 regression하는 네트워크 내부의 feature를 함께 concat해서 집어 넣어 주니 학습이 더 잘 된다고 합니다. 생각해보면 embedding vector라는 것이 우리가 쉽게 생각하던 것도 아니고 loss도 단순히 각도만 생각하기 때문에 학습이 쉬울 거 같지는 않습니다. 사실 이 연구진들이 기존에 아래 그림의 (a)와 같이 distance와 embedding을 각각 만드는 방식으로 학습을 진행했었는데 잘 안되는 케이스들이 있었다고 합니다.

지금 설명하고 있는 논문은 (b)의 형태를 띄고 있는데 D-feat을 함께 임베딩 벡터를 만들 때 사용하니 아래와 같이 모호한 케이스에 대해서 더 잘 된다고 합니다. 생각해보면 D-feat은 결국 boundary detection과 비슷한 역할을 하기 때문에 아래 사진에 대해서 긍정적인 영향을 줄 수 있을 것입니다.

 

디테일

또 한 가지 디테일로는 Embedding loss를 계산할 때 모든 인스턴스와 cosine similarity를 구하는 것이 아니라 가까운 애들끼리 loss를 계산하도록 합니다. 이는 위에서 설명드린 angular clustering이 어차피 seed point를 기준으로 진행되기 때문에 인접한 애들끼리만 다른 각도로 embedding vector가 형성되면 되기 때문입니다. 또한 만약 instance의 개수가 정말 많아진다면 embedding vector의 차원도 굉장히 커져야 하는데 그러면 메모리나 computation이 정말 많이 필요하고, 비효율적일 것입니다. 그래서 얼마의 거리 이내에 있는 instance 끼리만 embedding loss를 계산하도록 되어 있습니다. 이렇게 하면 4색 정리와 같이 4개의 차원으로도 효율적으로 임베딩 벡터를 표현할 수 있을 것입니다. 그래서 이들은 임베딩 벡터의 차원을 4, 8, 16, 32, 64로도 해보고 인접한 instance만 따질지 전체 다 따질지 등의 ablation study를 진행했습니다.

 

결론

일단 이들이 Leaf segmentation challenge에서 몇 개의 분야에서 1위를 차지하였고 잘 작동한다고 합니다.

픽셀 임베딩 벡터 라는 개념과 cosine similarity 를 활용한 loss function 이 인상깊었습니다.