이 게시물에서는 밑의 논문에서 소개한 MaskFormer에 대해서 공부해보도록 하겠다.
https://arxiv.org/abs/2107.06278
Per-Pixel Classification is Not All You Need for Semantic Segmentation
Modern approaches typically formulate semantic segmentation as a per-pixel classification task, while instance-level segmentation is handled with an alternative mask classification. Our key insight: mask classification is sufficiently general to solve both
arxiv.org
MaskFormer
MaskFormer의 목적은 instance segmentation과 semantic segmentation을 동시에 수행할 수 있는 model을 만드는 것이다.
Method
전체 알고리즘 개요도는 밑의 그림과 같다.
Pixel-level module에서는 먼저 이미지가 U-Net 기반 FCN(backbone) 을 통과하여 채널 수가 C_F이고, 공간 해상도는 H/S × W/S인 feature map을 생성한다. 이후 이 feature map을 Transformer decoder로 전달한다.
Transformer decoder의 입력은 학습 가능한 object query들이다. 각 query는 “하나의 object를 나타내는 latent slot”으로 동작한다고 이해할 수 있다. Decoder에서는 먼저 query들끼리 self-attention을 수행해, 각 query가 서로 다른 object를 담당하도록 역할을 조정한다. 그 후 feature map의 token들에 대해 cross-attention을 수행하면서, 각 query는 자신이 예측해야 할 object와 관련된 이미지 영역의 정보를 받아온다. 이러한 과정을 통해 각 query는 단일 object에 대응되는 embedding을 얻게 된다.
Decoder 출력을 바탕으로, 각 query embedding은 두 개의 독립된 MLP에 입력되어 class probability vector와 mask embedding을 생성한다. 이제 Class는 바로 예측할 수 있으므로, 이후에는 mask 생성 과정만 수행하면 된다.
Mask 생성은 다음과 같다. Pixel decoder를 통해 얻은 채널 수 C_e이고 해상도가 H × W인 feature map과, 각 query의 mask embedding(C_e 차원) 을 pixel-wise dot product 하면, 각 query마다 H × W 해상도의 mask probability map이 만들어진다. 이제 N개의 (class,mask) output이 준비됐다.
훈련 과정에서는 보통 GT 이미지에 K개의 object mask(K < N) 가 존재한다. 모델은 총 N개의 query를 출력하므로, Hungarian matching을 사용해 N개의 query와 K개의 GT object를 매칭한다. 매칭되지 않은 N-K개의 query는 class loss만(no-object로) 적용되고, 매칭된 query들은 class loss + mask loss를 모두 받게 된다.
inference시에는 semantic segmentation의 경우 각 픽셀별로 모든 mask를 돌면서 각 class별로 class prob x mask prob을 sum해서 class별 확률을 만들고 그것의 argmax를 취한다. 그리고 instance segmentation의 경우에는 각 픽셀별로 모든 mask중에서 max class prob x mask prob이 최대인 mask를 골라서 class, mask index를 기록하는 방식으로 동작한다.